Skip to main content

mongreldb_query/
extended_sql_functions.rs

1//! Extended SQL Functions: application-oriented scalar functions layered on
2//! DataFusion for date/time, JSON, string, and math compatibility.
3
4use arrow::array::{ArrayRef, BinaryArray, Float64Array, Int64Array, StringArray};
5use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6use arrow::record_batch::RecordBatch;
7use chrono::{
8    DateTime, Datelike, Duration, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Timelike, Utc,
9};
10use datafusion::catalog::{TableFunctionArgs, TableFunctionImpl, TableProvider};
11use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue};
12use datafusion::datasource::MemTable;
13use datafusion::logical_expr::{
14    ColumnarValue, Expr as DfExpr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature,
15    TypeSignature, Volatility,
16};
17use datafusion::prelude::SessionContext;
18use serde_json::{Map as JsonMap, Number as JsonNumber, Value as JsonValue};
19use std::hash::{Hash, Hasher};
20use std::sync::atomic::{AtomicU64, Ordering};
21use std::sync::Arc;
22
23static RANDOM_STATE: AtomicU64 = AtomicU64::new(0x9E37_79B9_7F4A_7C15);
24
25#[derive(Debug, Default)]
26pub struct ExtendedSqlState {
27    last_changes: AtomicU64,
28    total_changes: AtomicU64,
29    last_insert_rowid: AtomicU64,
30}
31
32impl ExtendedSqlState {
33    pub fn record_changes(&self, changes: u64, last_insert_rowid: Option<u64>) {
34        self.last_changes.store(changes, Ordering::Relaxed);
35        self.total_changes.fetch_add(changes, Ordering::Relaxed);
36        if let Some(rowid) = last_insert_rowid {
37            self.last_insert_rowid.store(rowid, Ordering::Relaxed);
38        }
39    }
40
41    fn last_changes(&self) -> u64 {
42        self.last_changes.load(Ordering::Relaxed)
43    }
44
45    fn total_changes(&self) -> u64 {
46        self.total_changes.load(Ordering::Relaxed)
47    }
48
49    fn last_insert_rowid(&self) -> u64 {
50        self.last_insert_rowid.load(Ordering::Relaxed)
51    }
52}
53
54pub fn register_extended_sql_functions(ctx: &SessionContext) {
55    register_extended_sql_functions_with_state(ctx, Arc::new(ExtendedSqlState::default()));
56}
57
58pub fn register_extended_sql_functions_with_state(
59    ctx: &SessionContext,
60    state: Arc<ExtendedSqlState>,
61) {
62    for func in extended_functions() {
63        ctx.register_udf(ScalarUDF::from(func.with_state(Arc::clone(&state))));
64    }
65    ctx.register_udtf("json_each", Arc::new(JsonEachFunc));
66    ctx.register_udtf("json_tree", Arc::new(JsonTreeFunc));
67    ctx.register_udtf("jsonb_each", Arc::new(JsonbEachFunc));
68    ctx.register_udtf("jsonb_tree", Arc::new(JsonbTreeFunc));
69    ctx.register_udtf("series", Arc::new(SeriesFunc));
70}
71
72pub fn extended_sql_function_names() -> Vec<&'static str> {
73    vec![
74        "date",
75        "time",
76        "datetime",
77        "julianday",
78        "unixepoch",
79        "strftime",
80        "timediff",
81        "json",
82        "jsonb",
83        "json_valid",
84        "json_extract",
85        "jsonb_extract",
86        "json_type",
87        "json_array_length",
88        "json_quote",
89        "json_array",
90        "jsonb_array",
91        "json_object",
92        "jsonb_object",
93        "json_patch",
94        "jsonb_patch",
95        "json_pretty",
96        "json_error_position",
97        "json_array_insert",
98        "jsonb_array_insert",
99        "json_set",
100        "jsonb_set",
101        "json_insert",
102        "jsonb_insert",
103        "json_replace",
104        "jsonb_replace",
105        "json_remove",
106        "jsonb_remove",
107        "json_each",
108        "json_tree",
109        "jsonb_each",
110        "jsonb_tree",
111        "series",
112        "abs",
113        "changes",
114        "instr",
115        "last_insert_rowid",
116        "quote",
117        "hex",
118        "unhex",
119        "printf",
120        "format",
121        "char",
122        "coalesce",
123        "concat",
124        "concat_ws",
125        "like",
126        "mod",
127        "pow",
128        "power",
129        "pi",
130        "random",
131        "randomblob",
132        "zeroblob",
133        "glob",
134        "length",
135        "octet_length",
136        "replace",
137        "round",
138        "sign",
139        "substr",
140        "substring",
141        "typeof",
142        "unicode",
143        "unistr",
144        "unistr_quote",
145        "lower",
146        "upper",
147        "ltrim",
148        "rtrim",
149        "trim",
150        "likely",
151        "unlikely",
152        "likelihood",
153        "ifnull",
154        "iif",
155        "if",
156        "nullif",
157        "soundex",
158        "sqlite_compileoption_get",
159        "sqlite_compileoption_used",
160        "sqlite_offset",
161        "sqlite_source_id",
162        "sqlite_version",
163        "total_changes",
164        "load_extension",
165        "acos",
166        "acosh",
167        "asin",
168        "asinh",
169        "atan",
170        "atan2",
171        "atanh",
172        "ceil",
173        "ceiling",
174        "cos",
175        "cosh",
176        "degrees",
177        "exp",
178        "floor",
179        "ln",
180        "log",
181        "log10",
182        "log2",
183        "radians",
184        "sin",
185        "sinh",
186        "sqrt",
187        "tan",
188        "tanh",
189    ]
190}
191
192pub fn contains_volatile_extended_function(sql: &str) -> bool {
193    let lower = sql.to_ascii_lowercase();
194    [
195        "date",
196        "time",
197        "datetime",
198        "julianday",
199        "unixepoch",
200        "strftime",
201        "timediff",
202        "random",
203        "randomblob",
204    ]
205    .iter()
206    .any(|name| contains_function_call(&lower, name))
207}
208
209fn contains_function_call(sql: &str, name: &str) -> bool {
210    let mut offset = 0;
211    while let Some(pos) = sql[offset..].find(name) {
212        let start = offset + pos;
213        let end = start + name.len();
214        let before_boundary = start == 0 || !is_ident_byte(sql.as_bytes()[start - 1]);
215        let after_boundary = sql
216            .as_bytes()
217            .get(end)
218            .map(|b| !is_ident_byte(*b))
219            .unwrap_or(true);
220        if before_boundary && after_boundary {
221            let rest = &sql[end..];
222            if rest.trim_start().starts_with('(') {
223                return true;
224            }
225        }
226        offset = end;
227    }
228    false
229}
230
231fn is_ident_byte(b: u8) -> bool {
232    b.is_ascii_alphanumeric() || b == b'_'
233}
234
235#[derive(Debug, Clone)]
236struct ExtendedSqlFunc {
237    name: &'static str,
238    aliases: Vec<String>,
239    kind: ExtendedFuncKind,
240    signature: Signature,
241    state: Arc<ExtendedSqlState>,
242}
243
244impl PartialEq for ExtendedSqlFunc {
245    fn eq(&self, other: &Self) -> bool {
246        self.name == other.name
247            && self.aliases == other.aliases
248            && self.kind == other.kind
249            && self.signature == other.signature
250    }
251}
252
253impl Eq for ExtendedSqlFunc {}
254
255impl Hash for ExtendedSqlFunc {
256    fn hash<H: Hasher>(&self, state: &mut H) {
257        self.name.hash(state);
258        self.aliases.hash(state);
259        self.kind.hash(state);
260        self.signature.hash(state);
261    }
262}
263
264impl ExtendedSqlFunc {
265    fn new(name: &'static str, kind: ExtendedFuncKind, volatility: Volatility) -> Self {
266        Self {
267            name,
268            aliases: Vec::new(),
269            kind,
270            signature: Signature::variadic_any(volatility),
271            state: Arc::new(ExtendedSqlState::default()),
272        }
273    }
274
275    fn zero_or_more(name: &'static str, kind: ExtendedFuncKind, volatility: Volatility) -> Self {
276        Self {
277            name,
278            aliases: Vec::new(),
279            kind,
280            signature: Signature::new(
281                TypeSignature::OneOf(vec![TypeSignature::Nullary, TypeSignature::VariadicAny]),
282                volatility,
283            ),
284            state: Arc::new(ExtendedSqlState::default()),
285        }
286    }
287
288    fn nullary(name: &'static str, kind: ExtendedFuncKind, volatility: Volatility) -> Self {
289        Self {
290            name,
291            aliases: Vec::new(),
292            kind,
293            signature: Signature::new(TypeSignature::Nullary, volatility),
294            state: Arc::new(ExtendedSqlState::default()),
295        }
296    }
297
298    fn with_aliases(mut self, aliases: &[&str]) -> Self {
299        self.aliases = aliases.iter().map(|s| s.to_string()).collect();
300        self
301    }
302
303    fn with_state(mut self, state: Arc<ExtendedSqlState>) -> Self {
304        self.state = state;
305        self
306    }
307}
308
309#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
310enum ExtendedFuncKind {
311    Date,
312    Time,
313    DateTime,
314    JulianDay,
315    UnixEpoch,
316    Strftime,
317    TimeDiff,
318    Json,
319    JsonValid,
320    JsonExtract,
321    JsonType,
322    JsonArrayLength,
323    JsonQuote,
324    JsonArray,
325    JsonObject,
326    JsonPatch,
327    JsonPretty,
328    JsonErrorPosition,
329    JsonArrayInsert,
330    JsonSet,
331    JsonInsert,
332    JsonReplace,
333    JsonRemove,
334    Abs,
335    Changes,
336    TotalChanges,
337    LastInsertRowid,
338    Instr,
339    Quote,
340    Hex,
341    Unhex,
342    Printf,
343    Char,
344    Coalesce,
345    Concat,
346    ConcatWs,
347    Like,
348    Mod,
349    Pow,
350    Pi,
351    Random,
352    RandomBlob,
353    ZeroBlob,
354    Glob,
355    Length,
356    OctetLength,
357    Replace,
358    Round,
359    Sign,
360    Substr,
361    TypeOf,
362    Unicode,
363    Unistr,
364    UnistrQuote,
365    Lower,
366    Upper,
367    Ltrim,
368    Rtrim,
369    Trim,
370    Likely,
371    Unlikely,
372    Likelihood,
373    IfNull,
374    Iif,
375    NullIf,
376    ScalarMax,
377    ScalarMin,
378    Soundex,
379    CompileOptionGet,
380    CompileOptionUsed,
381    Offset,
382    SourceId,
383    Version,
384    LoadExtension,
385    Acos,
386    Acosh,
387    Asin,
388    Asinh,
389    Atan,
390    Atan2,
391    Atanh,
392    Ceil,
393    Cos,
394    Cosh,
395    Degrees,
396    Exp,
397    Floor,
398    Ln,
399    Log,
400    Log10,
401    Log2,
402    Radians,
403    Sin,
404    Sinh,
405    Sqrt,
406    Tan,
407    Tanh,
408}
409
410impl ScalarUDFImpl for ExtendedSqlFunc {
411    fn name(&self) -> &str {
412        self.name
413    }
414
415    fn aliases(&self) -> &[String] {
416        &self.aliases
417    }
418
419    fn signature(&self) -> &Signature {
420        &self.signature
421    }
422
423    fn return_type(&self, _args: &[DataType]) -> DFResult<DataType> {
424        Ok(match self.kind {
425            ExtendedFuncKind::JsonValid
426            | ExtendedFuncKind::JsonArrayLength
427            | ExtendedFuncKind::Instr
428            | ExtendedFuncKind::UnixEpoch
429            | ExtendedFuncKind::Random
430            | ExtendedFuncKind::Changes
431            | ExtendedFuncKind::TotalChanges
432            | ExtendedFuncKind::LastInsertRowid
433            | ExtendedFuncKind::Glob
434            | ExtendedFuncKind::Like
435            | ExtendedFuncKind::Length
436            | ExtendedFuncKind::OctetLength
437            | ExtendedFuncKind::Sign
438            | ExtendedFuncKind::Unicode
439            | ExtendedFuncKind::CompileOptionUsed
440            | ExtendedFuncKind::Offset
441            | ExtendedFuncKind::JsonErrorPosition => DataType::Int64,
442            ExtendedFuncKind::JulianDay
443            | ExtendedFuncKind::TimeDiff
444            | ExtendedFuncKind::Mod
445            | ExtendedFuncKind::Pow
446            | ExtendedFuncKind::Pi
447            | ExtendedFuncKind::Round
448            | ExtendedFuncKind::Acos
449            | ExtendedFuncKind::Acosh
450            | ExtendedFuncKind::Asin
451            | ExtendedFuncKind::Asinh
452            | ExtendedFuncKind::Atan
453            | ExtendedFuncKind::Atan2
454            | ExtendedFuncKind::Atanh
455            | ExtendedFuncKind::Ceil
456            | ExtendedFuncKind::Cos
457            | ExtendedFuncKind::Cosh
458            | ExtendedFuncKind::Degrees
459            | ExtendedFuncKind::Exp
460            | ExtendedFuncKind::Floor
461            | ExtendedFuncKind::Ln
462            | ExtendedFuncKind::Log
463            | ExtendedFuncKind::Log10
464            | ExtendedFuncKind::Log2
465            | ExtendedFuncKind::Radians
466            | ExtendedFuncKind::Sin
467            | ExtendedFuncKind::Sinh
468            | ExtendedFuncKind::Sqrt
469            | ExtendedFuncKind::Tan
470            | ExtendedFuncKind::Tanh => DataType::Float64,
471            ExtendedFuncKind::RandomBlob | ExtendedFuncKind::ZeroBlob => DataType::Binary,
472            ExtendedFuncKind::Abs => _args.first().cloned().unwrap_or(DataType::Float64),
473            ExtendedFuncKind::Likely
474            | ExtendedFuncKind::Unlikely
475            | ExtendedFuncKind::Likelihood => _args.first().cloned().unwrap_or(DataType::Utf8),
476            ExtendedFuncKind::Coalesce => _args
477                .iter()
478                .find(|ty| **ty != DataType::Null)
479                .cloned()
480                .unwrap_or(DataType::Utf8),
481            ExtendedFuncKind::ScalarMax | ExtendedFuncKind::ScalarMin => _args
482                .iter()
483                .find(|ty| **ty != DataType::Null)
484                .cloned()
485                .unwrap_or(DataType::Utf8),
486            ExtendedFuncKind::IfNull => _args
487                .first()
488                .filter(|ty| **ty != DataType::Null)
489                .or_else(|| _args.get(1))
490                .cloned()
491                .unwrap_or(DataType::Utf8),
492            ExtendedFuncKind::Iif => _args.get(1).cloned().unwrap_or(DataType::Utf8),
493            ExtendedFuncKind::NullIf => _args.first().cloned().unwrap_or(DataType::Utf8),
494            _ => DataType::Utf8,
495        })
496    }
497
498    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DFResult<ColumnarValue> {
499        match self.kind {
500            ExtendedFuncKind::Date => eval_string(args, |row| {
501                parse_datetime_arg(row, 0).map(|dt| dt.format("%Y-%m-%d").to_string())
502            }),
503            ExtendedFuncKind::Time => eval_string(args, |row| {
504                parse_datetime_arg(row, 0).map(|dt| dt.format("%H:%M:%S").to_string())
505            }),
506            ExtendedFuncKind::DateTime => eval_string(args, |row| {
507                parse_datetime_arg(row, 0).map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
508            }),
509            ExtendedFuncKind::JulianDay => eval_float(args, |row| {
510                parse_datetime_arg(row, 0).map(|dt| unix_to_julian(dt.timestamp() as f64))
511            }),
512            ExtendedFuncKind::UnixEpoch => eval_int(args, |row| {
513                parse_datetime_arg(row, 0).map(|dt| dt.timestamp())
514            }),
515            ExtendedFuncKind::Strftime => eval_string(args, strftime_row),
516            ExtendedFuncKind::TimeDiff => eval_float(args, |row| {
517                let lhs = parse_datetime_at(row, 0)?;
518                let rhs = parse_datetime_at(row, 1)?;
519                Some((lhs.timestamp_millis() - rhs.timestamp_millis()) as f64 / 1000.0)
520            }),
521            ExtendedFuncKind::Json => eval_string(args, |row| {
522                let parsed = parse_json_arg(row, 0)?;
523                Some(parsed.to_string())
524            }),
525            ExtendedFuncKind::JsonValid => eval_int(args, |row| {
526                Some(
527                    if row
528                        .text(0)
529                        .as_deref()
530                        .and_then(|s| serde_json::from_str::<JsonValue>(s).ok())
531                        .is_some()
532                    {
533                        1
534                    } else {
535                        0
536                    },
537                )
538            }),
539            ExtendedFuncKind::JsonExtract => eval_string(args, json_extract_row),
540            ExtendedFuncKind::JsonType => eval_string(args, json_type_row),
541            ExtendedFuncKind::JsonArrayLength => eval_int(args, json_array_length_row),
542            ExtendedFuncKind::JsonQuote => eval_string(args, |row| {
543                row.scalar_any(0)
544                    .map(scalar_to_json_arg)
545                    .map(|v| v.to_string())
546            }),
547            ExtendedFuncKind::JsonArray => eval_string(args, json_array_row),
548            ExtendedFuncKind::JsonObject => eval_string(args, json_object_row),
549            ExtendedFuncKind::JsonPatch => eval_string(args, json_patch_row),
550            ExtendedFuncKind::JsonPretty => eval_string(args, |row| {
551                let parsed = parse_json_arg(row, 0)?;
552                serde_json::to_string_pretty(&parsed).ok()
553            }),
554            ExtendedFuncKind::JsonErrorPosition => eval_int(args, |row| {
555                Some(json_error_position(row.text(0).as_deref().unwrap_or_default()) as i64)
556            }),
557            ExtendedFuncKind::JsonArrayInsert => eval_string(args, json_array_insert_row),
558            ExtendedFuncKind::JsonSet => {
559                eval_string(args, |row| json_mutate_row(row, JsonMutation::Set))
560            }
561            ExtendedFuncKind::JsonInsert => {
562                eval_string(args, |row| json_mutate_row(row, JsonMutation::Insert))
563            }
564            ExtendedFuncKind::JsonReplace => {
565                eval_string(args, |row| json_mutate_row(row, JsonMutation::Replace))
566            }
567            ExtendedFuncKind::JsonRemove => eval_string(args, json_remove_row),
568            ExtendedFuncKind::Abs => eval_scalar(args, abs_scalar),
569            ExtendedFuncKind::Changes => {
570                let state = Arc::clone(&self.state);
571                eval_int(args, move |_| Some(state.last_changes() as i64))
572            }
573            ExtendedFuncKind::TotalChanges => {
574                let state = Arc::clone(&self.state);
575                eval_int(args, move |_| Some(state.total_changes() as i64))
576            }
577            ExtendedFuncKind::LastInsertRowid => {
578                let state = Arc::clone(&self.state);
579                eval_int(args, move |_| Some(state.last_insert_rowid() as i64))
580            }
581            ExtendedFuncKind::Instr => eval_int(args, |row| {
582                let haystack = row.text(0)?;
583                let needle = row.text(1)?;
584                Some(instr(&haystack, &needle))
585            }),
586            ExtendedFuncKind::Quote => eval_string(args, |row| row.scalar_any(0).map(sql_quote)),
587            ExtendedFuncKind::Hex => eval_string(args, |row| {
588                row.scalar_any(0).map(|v| hex_encode(&scalar_bytes(&v)))
589            }),
590            ExtendedFuncKind::Unhex => eval_string(args, |row| {
591                let ignored = row.text(1).unwrap_or_default();
592                let bytes = hex_decode_ignoring(&row.text(0)?, &ignored)?;
593                String::from_utf8(bytes).ok()
594            }),
595            ExtendedFuncKind::Printf => eval_string(args, printf_row),
596            ExtendedFuncKind::Char => eval_string(args, char_row),
597            ExtendedFuncKind::Coalesce => eval_scalar(args, coalesce_scalar),
598            ExtendedFuncKind::Concat => eval_string(args, concat_row),
599            ExtendedFuncKind::ConcatWs => eval_string(args, concat_ws_row),
600            ExtendedFuncKind::Like => eval_int(args, |row| {
601                let pattern = row.text(0)?;
602                let value = row.text(1)?;
603                let escape = row.text(2).and_then(|s| s.chars().next());
604                Some(i64::from(like_match_sql(&value, &pattern, escape)))
605            }),
606            ExtendedFuncKind::Mod => eval_float(args, |row| {
607                let lhs = row.f64(0)?;
608                let rhs = row.f64(1)?;
609                if rhs == 0.0 {
610                    None
611                } else {
612                    Some(lhs % rhs)
613                }
614            }),
615            ExtendedFuncKind::Pow => eval_float(args, |row| Some(row.f64(0)?.powf(row.f64(1)?))),
616            ExtendedFuncKind::Pi => eval_float(args, |_| Some(std::f64::consts::PI)),
617            ExtendedFuncKind::Random => eval_int(args, |_| Some(next_random())),
618            ExtendedFuncKind::RandomBlob => eval_binary(args, |row| {
619                let len = row.i64(0).unwrap_or(1).max(1) as usize;
620                Some(random_bytes(len))
621            }),
622            ExtendedFuncKind::ZeroBlob => eval_binary(args, |row| {
623                let len = row.i64(0).unwrap_or(1).max(1) as usize;
624                Some(vec![0; len])
625            }),
626            ExtendedFuncKind::Glob => eval_int(args, |row| {
627                let pattern = row.text(0)?;
628                let value = row.text(1)?;
629                Some(if glob_match(&pattern, &value) { 1 } else { 0 })
630            }),
631            ExtendedFuncKind::Length => eval_int(args, |row| {
632                let value = row.scalar(0)?;
633                Some(match value {
634                    ScalarValue::Binary(Some(v))
635                    | ScalarValue::LargeBinary(Some(v))
636                    | ScalarValue::BinaryView(Some(v)) => v.len() as i64,
637                    _ => scalar_text(&value).chars().count() as i64,
638                })
639            }),
640            ExtendedFuncKind::OctetLength => eval_int(args, |row| {
641                let value = row.scalar(0)?;
642                Some(scalar_bytes(&value).len() as i64)
643            }),
644            ExtendedFuncKind::Replace => eval_string(args, |row| {
645                let input = row.text(0)?;
646                let from = row.text(1)?;
647                let to = row.text(2)?;
648                if from.is_empty() {
649                    Some(input)
650                } else {
651                    Some(input.replace(&from, &to))
652                }
653            }),
654            ExtendedFuncKind::Round => eval_float(args, |row| {
655                let value = row.f64(0)?;
656                let places = row.i64(1).unwrap_or(0).max(0) as i32;
657                let scale = 10_f64.powi(places);
658                Some((value * scale).round() / scale)
659            }),
660            ExtendedFuncKind::Sign => eval_int(args, |row| {
661                let value = row.f64(0)?;
662                Some(if value > 0.0 {
663                    1
664                } else if value < 0.0 {
665                    -1
666                } else {
667                    0
668                })
669            }),
670            ExtendedFuncKind::Substr => eval_string(args, substr_row),
671            ExtendedFuncKind::TypeOf => eval_string(args, |row| {
672                row.scalar_any(0).map(|v| scalar_type_name(&v).to_string())
673            }),
674            ExtendedFuncKind::Unicode => {
675                eval_int(args, |row| row.text(0)?.chars().next().map(|ch| ch as i64))
676            }
677            ExtendedFuncKind::Unistr => eval_string(args, |row| unistr(&row.text(0)?)),
678            ExtendedFuncKind::UnistrQuote => eval_string(args, |row| {
679                row.text(0)
680                    .map(|value| sql_quote(ScalarValue::Utf8(Some(unistr_quote_text(&value)))))
681            }),
682            ExtendedFuncKind::Lower => eval_string(args, |row| Some(row.text(0)?.to_lowercase())),
683            ExtendedFuncKind::Upper => eval_string(args, |row| Some(row.text(0)?.to_uppercase())),
684            ExtendedFuncKind::Ltrim => eval_string(args, |row| {
685                let chars = row.text(1).unwrap_or_else(|| " ".to_string());
686                Some(
687                    row.text(0)?
688                        .trim_start_matches(|c| chars.contains(c))
689                        .to_string(),
690                )
691            }),
692            ExtendedFuncKind::Rtrim => eval_string(args, |row| {
693                let chars = row.text(1).unwrap_or_else(|| " ".to_string());
694                Some(
695                    row.text(0)?
696                        .trim_end_matches(|c| chars.contains(c))
697                        .to_string(),
698                )
699            }),
700            ExtendedFuncKind::Trim => eval_string(args, |row| {
701                let chars = row.text(1).unwrap_or_else(|| " ".to_string());
702                Some(row.text(0)?.trim_matches(|c| chars.contains(c)).to_string())
703            }),
704            ExtendedFuncKind::Likely | ExtendedFuncKind::Unlikely => {
705                eval_scalar(args, |row| row.scalar_any(0).unwrap_or(ScalarValue::Null))
706            }
707            ExtendedFuncKind::Likelihood => {
708                eval_scalar(args, |row| row.scalar_any(0).unwrap_or(ScalarValue::Null))
709            }
710            ExtendedFuncKind::IfNull => eval_scalar(args, |row| {
711                let first = row.scalar_any(0).unwrap_or(ScalarValue::Null);
712                if first.is_null() {
713                    row.scalar_any(1).unwrap_or(ScalarValue::Null)
714                } else {
715                    first
716                }
717            }),
718            ExtendedFuncKind::Iif => eval_scalar(args, |row| {
719                if row.truthy(0) {
720                    row.scalar_any(1).unwrap_or(ScalarValue::Null)
721                } else {
722                    row.scalar_any(2).unwrap_or(ScalarValue::Null)
723                }
724            }),
725            ExtendedFuncKind::NullIf => eval_scalar(args, |row| {
726                let first = row.scalar_any(0).unwrap_or(ScalarValue::Null);
727                let first_type = first.data_type();
728                let second = row.scalar_any(1).unwrap_or(ScalarValue::Null);
729                if !first.is_null()
730                    && !second.is_null()
731                    && scalar_text(&first) == scalar_text(&second)
732                {
733                    ScalarValue::try_new_null(&first_type).unwrap_or(ScalarValue::Null)
734                } else {
735                    first
736                }
737            }),
738            ExtendedFuncKind::ScalarMax => eval_scalar(args, |row| scalar_extreme(row, true)),
739            ExtendedFuncKind::ScalarMin => eval_scalar(args, |row| scalar_extreme(row, false)),
740            ExtendedFuncKind::Soundex => eval_string(args, |row| Some(soundex(&row.text(0)?))),
741            ExtendedFuncKind::CompileOptionGet => eval_string(args, |row| {
742                row.i64(0)
743                    .and_then(|idx| compile_options().get(idx as usize).map(|s| s.to_string()))
744            }),
745            ExtendedFuncKind::CompileOptionUsed => eval_int(args, |row| {
746                let needle = row.text(0)?.to_ascii_uppercase();
747                Some(i64::from(
748                    compile_options()
749                        .iter()
750                        .any(|opt| opt.eq_ignore_ascii_case(&needle)),
751                ))
752            }),
753            ExtendedFuncKind::Offset => eval_int(args, |_| None),
754            ExtendedFuncKind::SourceId => eval_string(args, |_| {
755                Some(format!(
756                    "mongreldb-query {} extended-sql-profile",
757                    env!("CARGO_PKG_VERSION")
758                ))
759            }),
760            ExtendedFuncKind::Version => {
761                eval_string(args, |_| Some(env!("CARGO_PKG_VERSION").to_string()))
762            }
763            ExtendedFuncKind::LoadExtension => Err(DataFusionError::Execution(
764                "load_extension is disabled in MongrelDB SQL sessions".into(),
765            )),
766            ExtendedFuncKind::Acos => eval_float(args, |row| math_unary(row, f64::acos)),
767            ExtendedFuncKind::Acosh => eval_float(args, |row| math_unary(row, f64::acosh)),
768            ExtendedFuncKind::Asin => eval_float(args, |row| math_unary(row, f64::asin)),
769            ExtendedFuncKind::Asinh => eval_float(args, |row| math_unary(row, f64::asinh)),
770            ExtendedFuncKind::Atan => eval_float(args, |row| math_unary(row, f64::atan)),
771            ExtendedFuncKind::Atan2 => eval_float(args, |row| Some(row.f64(0)?.atan2(row.f64(1)?))),
772            ExtendedFuncKind::Atanh => eval_float(args, |row| math_unary(row, f64::atanh)),
773            ExtendedFuncKind::Ceil => eval_float(args, |row| math_unary(row, f64::ceil)),
774            ExtendedFuncKind::Cos => eval_float(args, |row| math_unary(row, f64::cos)),
775            ExtendedFuncKind::Cosh => eval_float(args, |row| math_unary(row, f64::cosh)),
776            ExtendedFuncKind::Degrees => eval_float(args, |row| math_unary(row, f64::to_degrees)),
777            ExtendedFuncKind::Exp => eval_float(args, |row| math_unary(row, f64::exp)),
778            ExtendedFuncKind::Floor => eval_float(args, |row| math_unary(row, f64::floor)),
779            ExtendedFuncKind::Ln => eval_float(args, |row| math_unary(row, f64::ln)),
780            ExtendedFuncKind::Log => eval_float(args, log_row),
781            ExtendedFuncKind::Log10 => eval_float(args, |row| math_unary(row, f64::log10)),
782            ExtendedFuncKind::Log2 => eval_float(args, |row| math_unary(row, f64::log2)),
783            ExtendedFuncKind::Radians => eval_float(args, |row| math_unary(row, f64::to_radians)),
784            ExtendedFuncKind::Sin => eval_float(args, |row| math_unary(row, f64::sin)),
785            ExtendedFuncKind::Sinh => eval_float(args, |row| math_unary(row, f64::sinh)),
786            ExtendedFuncKind::Sqrt => eval_float(args, |row| math_unary(row, f64::sqrt)),
787            ExtendedFuncKind::Tan => eval_float(args, |row| math_unary(row, f64::tan)),
788            ExtendedFuncKind::Tanh => eval_float(args, |row| math_unary(row, f64::tanh)),
789        }
790    }
791}
792
793fn extended_functions() -> Vec<ExtendedSqlFunc> {
794    use ExtendedFuncKind::*;
795    vec![
796        ExtendedSqlFunc::zero_or_more("date", Date, Volatility::Volatile),
797        ExtendedSqlFunc::zero_or_more("time", Time, Volatility::Volatile),
798        ExtendedSqlFunc::zero_or_more("datetime", DateTime, Volatility::Volatile),
799        ExtendedSqlFunc::zero_or_more("julianday", JulianDay, Volatility::Volatile),
800        ExtendedSqlFunc::zero_or_more("unixepoch", UnixEpoch, Volatility::Volatile),
801        ExtendedSqlFunc::new("strftime", Strftime, Volatility::Volatile),
802        ExtendedSqlFunc::new("timediff", TimeDiff, Volatility::Volatile),
803        ExtendedSqlFunc::new("json", Json, Volatility::Immutable).with_aliases(&["jsonb"]),
804        ExtendedSqlFunc::new("json_valid", JsonValid, Volatility::Immutable),
805        ExtendedSqlFunc::new("json_extract", JsonExtract, Volatility::Immutable)
806            .with_aliases(&["jsonb_extract"]),
807        ExtendedSqlFunc::new("json_type", JsonType, Volatility::Immutable),
808        ExtendedSqlFunc::new("json_array_length", JsonArrayLength, Volatility::Immutable),
809        ExtendedSqlFunc::new("json_quote", JsonQuote, Volatility::Immutable),
810        ExtendedSqlFunc::new("json_array", JsonArray, Volatility::Immutable)
811            .with_aliases(&["jsonb_array"]),
812        ExtendedSqlFunc::new("json_object", JsonObject, Volatility::Immutable)
813            .with_aliases(&["jsonb_object"]),
814        ExtendedSqlFunc::new("json_patch", JsonPatch, Volatility::Immutable)
815            .with_aliases(&["jsonb_patch"]),
816        ExtendedSqlFunc::new("json_pretty", JsonPretty, Volatility::Immutable),
817        ExtendedSqlFunc::new(
818            "json_error_position",
819            JsonErrorPosition,
820            Volatility::Immutable,
821        ),
822        ExtendedSqlFunc::new("json_array_insert", JsonArrayInsert, Volatility::Immutable)
823            .with_aliases(&["jsonb_array_insert"]),
824        ExtendedSqlFunc::new("json_set", JsonSet, Volatility::Immutable)
825            .with_aliases(&["jsonb_set"]),
826        ExtendedSqlFunc::new("json_insert", JsonInsert, Volatility::Immutable)
827            .with_aliases(&["jsonb_insert"]),
828        ExtendedSqlFunc::new("json_replace", JsonReplace, Volatility::Immutable)
829            .with_aliases(&["jsonb_replace"]),
830        ExtendedSqlFunc::new("json_remove", JsonRemove, Volatility::Immutable)
831            .with_aliases(&["jsonb_remove"]),
832        ExtendedSqlFunc::new("abs", Abs, Volatility::Immutable),
833        ExtendedSqlFunc::nullary("changes", Changes, Volatility::Volatile),
834        ExtendedSqlFunc::nullary("total_changes", TotalChanges, Volatility::Volatile),
835        ExtendedSqlFunc::nullary("last_insert_rowid", LastInsertRowid, Volatility::Volatile),
836        ExtendedSqlFunc::new("instr", Instr, Volatility::Immutable),
837        ExtendedSqlFunc::new("quote", Quote, Volatility::Immutable),
838        ExtendedSqlFunc::new("hex", Hex, Volatility::Immutable),
839        ExtendedSqlFunc::new("unhex", Unhex, Volatility::Immutable),
840        ExtendedSqlFunc::new("printf", Printf, Volatility::Immutable).with_aliases(&["format"]),
841        ExtendedSqlFunc::new("char", Char, Volatility::Immutable),
842        ExtendedSqlFunc::new("coalesce", Coalesce, Volatility::Immutable),
843        ExtendedSqlFunc::new("concat", Concat, Volatility::Immutable),
844        ExtendedSqlFunc::new("concat_ws", ConcatWs, Volatility::Immutable),
845        ExtendedSqlFunc::new("like", Like, Volatility::Immutable),
846        ExtendedSqlFunc::new("mod", Mod, Volatility::Immutable),
847        ExtendedSqlFunc::new("pow", Pow, Volatility::Immutable).with_aliases(&["power"]),
848        ExtendedSqlFunc::nullary("pi", Pi, Volatility::Immutable),
849        ExtendedSqlFunc::nullary("random", Random, Volatility::Volatile),
850        ExtendedSqlFunc::new("randomblob", RandomBlob, Volatility::Volatile),
851        ExtendedSqlFunc::new("zeroblob", ZeroBlob, Volatility::Immutable),
852        ExtendedSqlFunc::new("glob", Glob, Volatility::Immutable),
853        ExtendedSqlFunc::new("length", Length, Volatility::Immutable),
854        ExtendedSqlFunc::new("octet_length", OctetLength, Volatility::Immutable),
855        ExtendedSqlFunc::new("replace", Replace, Volatility::Immutable),
856        ExtendedSqlFunc::new("round", Round, Volatility::Immutable),
857        ExtendedSqlFunc::new("sign", Sign, Volatility::Immutable),
858        ExtendedSqlFunc::new("substr", Substr, Volatility::Immutable).with_aliases(&["substring"]),
859        ExtendedSqlFunc::new("typeof", TypeOf, Volatility::Immutable),
860        ExtendedSqlFunc::new("unicode", Unicode, Volatility::Immutable),
861        ExtendedSqlFunc::new("unistr", Unistr, Volatility::Immutable),
862        ExtendedSqlFunc::new("unistr_quote", UnistrQuote, Volatility::Immutable),
863        ExtendedSqlFunc::new("lower", Lower, Volatility::Immutable),
864        ExtendedSqlFunc::new("upper", Upper, Volatility::Immutable),
865        ExtendedSqlFunc::new("ltrim", Ltrim, Volatility::Immutable),
866        ExtendedSqlFunc::new("rtrim", Rtrim, Volatility::Immutable),
867        ExtendedSqlFunc::new("trim", Trim, Volatility::Immutable),
868        ExtendedSqlFunc::new("likely", Likely, Volatility::Immutable),
869        ExtendedSqlFunc::new("unlikely", Unlikely, Volatility::Immutable),
870        ExtendedSqlFunc::new("likelihood", Likelihood, Volatility::Immutable),
871        ExtendedSqlFunc::new("ifnull", IfNull, Volatility::Immutable),
872        ExtendedSqlFunc::new("iif", Iif, Volatility::Immutable).with_aliases(&["if"]),
873        ExtendedSqlFunc::new("nullif", NullIf, Volatility::Immutable),
874        ExtendedSqlFunc::new("__mongreldb_scalar_max", ScalarMax, Volatility::Immutable),
875        ExtendedSqlFunc::new("__mongreldb_scalar_min", ScalarMin, Volatility::Immutable),
876        ExtendedSqlFunc::new("soundex", Soundex, Volatility::Immutable),
877        ExtendedSqlFunc::new(
878            "sqlite_compileoption_get",
879            CompileOptionGet,
880            Volatility::Immutable,
881        ),
882        ExtendedSqlFunc::new(
883            "sqlite_compileoption_used",
884            CompileOptionUsed,
885            Volatility::Immutable,
886        ),
887        ExtendedSqlFunc::new("sqlite_offset", Offset, Volatility::Immutable),
888        ExtendedSqlFunc::nullary("sqlite_source_id", SourceId, Volatility::Immutable),
889        ExtendedSqlFunc::nullary("sqlite_version", Version, Volatility::Immutable),
890        ExtendedSqlFunc::new("load_extension", LoadExtension, Volatility::Volatile),
891        ExtendedSqlFunc::new("acos", Acos, Volatility::Immutable),
892        ExtendedSqlFunc::new("acosh", Acosh, Volatility::Immutable),
893        ExtendedSqlFunc::new("asin", Asin, Volatility::Immutable),
894        ExtendedSqlFunc::new("asinh", Asinh, Volatility::Immutable),
895        ExtendedSqlFunc::new("atan", Atan, Volatility::Immutable),
896        ExtendedSqlFunc::new("atan2", Atan2, Volatility::Immutable),
897        ExtendedSqlFunc::new("atanh", Atanh, Volatility::Immutable),
898        ExtendedSqlFunc::new("ceil", Ceil, Volatility::Immutable).with_aliases(&["ceiling"]),
899        ExtendedSqlFunc::new("cos", Cos, Volatility::Immutable),
900        ExtendedSqlFunc::new("cosh", Cosh, Volatility::Immutable),
901        ExtendedSqlFunc::new("degrees", Degrees, Volatility::Immutable),
902        ExtendedSqlFunc::new("exp", Exp, Volatility::Immutable),
903        ExtendedSqlFunc::new("floor", Floor, Volatility::Immutable),
904        ExtendedSqlFunc::new("ln", Ln, Volatility::Immutable),
905        ExtendedSqlFunc::new("log", Log, Volatility::Immutable),
906        ExtendedSqlFunc::new("log10", Log10, Volatility::Immutable),
907        ExtendedSqlFunc::new("log2", Log2, Volatility::Immutable),
908        ExtendedSqlFunc::new("radians", Radians, Volatility::Immutable),
909        ExtendedSqlFunc::new("sin", Sin, Volatility::Immutable),
910        ExtendedSqlFunc::new("sinh", Sinh, Volatility::Immutable),
911        ExtendedSqlFunc::new("sqrt", Sqrt, Volatility::Immutable),
912        ExtendedSqlFunc::new("tan", Tan, Volatility::Immutable),
913        ExtendedSqlFunc::new("tanh", Tanh, Volatility::Immutable),
914    ]
915}
916
917struct RowArgs {
918    arrays: Vec<ArrayRef>,
919    row: usize,
920}
921
922impl RowArgs {
923    fn scalar(&self, idx: usize) -> Option<ScalarValue> {
924        self.scalar_any(idx).filter(|v| !v.is_null())
925    }
926
927    fn scalar_any(&self, idx: usize) -> Option<ScalarValue> {
928        self.arrays
929            .get(idx)
930            .and_then(|arr| ScalarValue::try_from_array(arr, self.row).ok())
931    }
932
933    fn text(&self, idx: usize) -> Option<String> {
934        self.scalar(idx).map(|v| scalar_text(&v))
935    }
936
937    fn i64(&self, idx: usize) -> Option<i64> {
938        scalar_i64(&self.scalar(idx)?)
939    }
940
941    fn f64(&self, idx: usize) -> Option<f64> {
942        scalar_f64(&self.scalar(idx)?)
943    }
944
945    fn len(&self) -> usize {
946        self.arrays.len()
947    }
948
949    fn truthy(&self, idx: usize) -> bool {
950        match self.scalar(idx) {
951            Some(ScalarValue::Boolean(Some(v))) => v,
952            Some(v) => scalar_f64(&v).is_some_and(|n| n != 0.0),
953            None => false,
954        }
955    }
956}
957
958fn eval_string(
959    args: ScalarFunctionArgs,
960    mut f: impl FnMut(&RowArgs) -> Option<String>,
961) -> DFResult<ColumnarValue> {
962    let (arrays, rows, scalar) = expand_args(args.args, args.number_rows)?;
963    let values: Vec<Option<String>> = (0..rows)
964        .map(|row| {
965            f(&RowArgs {
966                arrays: arrays.clone(),
967                row,
968            })
969        })
970        .collect();
971    finish(scalar, Arc::new(StringArray::from(values)))
972}
973
974fn eval_int(
975    args: ScalarFunctionArgs,
976    mut f: impl FnMut(&RowArgs) -> Option<i64>,
977) -> DFResult<ColumnarValue> {
978    let (arrays, rows, scalar) = expand_args(args.args, args.number_rows)?;
979    let values: Vec<Option<i64>> = (0..rows)
980        .map(|row| {
981            f(&RowArgs {
982                arrays: arrays.clone(),
983                row,
984            })
985        })
986        .collect();
987    finish(scalar, Arc::new(Int64Array::from(values)))
988}
989
990fn eval_float(
991    args: ScalarFunctionArgs,
992    mut f: impl FnMut(&RowArgs) -> Option<f64>,
993) -> DFResult<ColumnarValue> {
994    let (arrays, rows, scalar) = expand_args(args.args, args.number_rows)?;
995    let values: Vec<Option<f64>> = (0..rows)
996        .map(|row| {
997            f(&RowArgs {
998                arrays: arrays.clone(),
999                row,
1000            })
1001        })
1002        .collect();
1003    finish(scalar, Arc::new(Float64Array::from(values)))
1004}
1005
1006fn eval_binary(
1007    args: ScalarFunctionArgs,
1008    mut f: impl FnMut(&RowArgs) -> Option<Vec<u8>>,
1009) -> DFResult<ColumnarValue> {
1010    let (arrays, rows, scalar) = expand_args(args.args, args.number_rows)?;
1011    let values: Vec<Option<Vec<u8>>> = (0..rows)
1012        .map(|row| {
1013            f(&RowArgs {
1014                arrays: arrays.clone(),
1015                row,
1016            })
1017        })
1018        .collect();
1019    let borrowed = values
1020        .iter()
1021        .map(|value| value.as_deref())
1022        .collect::<Vec<_>>();
1023    finish(scalar, Arc::new(BinaryArray::from(borrowed)))
1024}
1025
1026fn eval_scalar(
1027    args: ScalarFunctionArgs,
1028    mut f: impl FnMut(&RowArgs) -> ScalarValue,
1029) -> DFResult<ColumnarValue> {
1030    let (arrays, rows, scalar) = expand_args(args.args, args.number_rows)?;
1031    let values: Vec<ScalarValue> = (0..rows)
1032        .map(|row| {
1033            f(&RowArgs {
1034                arrays: arrays.clone(),
1035                row,
1036            })
1037        })
1038        .collect();
1039    if scalar {
1040        Ok(ColumnarValue::Scalar(
1041            values.into_iter().next().unwrap_or(ScalarValue::Null),
1042        ))
1043    } else {
1044        Ok(ColumnarValue::Array(ScalarValue::iter_to_array(values)?))
1045    }
1046}
1047
1048fn expand_args(
1049    args: Vec<ColumnarValue>,
1050    number_rows: usize,
1051) -> DFResult<(Vec<ArrayRef>, usize, bool)> {
1052    let scalar = args
1053        .iter()
1054        .all(|arg| matches!(arg, ColumnarValue::Scalar(_)));
1055    let rows = if scalar {
1056        number_rows.max(1)
1057    } else {
1058        number_rows
1059    };
1060    let arrays = args
1061        .into_iter()
1062        .map(|arg| arg.into_array(rows))
1063        .collect::<DFResult<Vec<_>>>()?;
1064    Ok((arrays, rows, scalar))
1065}
1066
1067fn finish(scalar: bool, arr: ArrayRef) -> DFResult<ColumnarValue> {
1068    if scalar {
1069        Ok(ColumnarValue::Scalar(ScalarValue::try_from_array(&arr, 0)?))
1070    } else {
1071        Ok(ColumnarValue::Array(arr))
1072    }
1073}
1074
1075fn scalar_text(v: &ScalarValue) -> String {
1076    match v {
1077        ScalarValue::Utf8(Some(s))
1078        | ScalarValue::LargeUtf8(Some(s))
1079        | ScalarValue::Utf8View(Some(s)) => s.clone(),
1080        ScalarValue::Boolean(Some(v)) => {
1081            if *v {
1082                "1".into()
1083            } else {
1084                "0".into()
1085            }
1086        }
1087        ScalarValue::Int8(Some(v)) => v.to_string(),
1088        ScalarValue::Int16(Some(v)) => v.to_string(),
1089        ScalarValue::Int32(Some(v)) => v.to_string(),
1090        ScalarValue::Int64(Some(v)) => v.to_string(),
1091        ScalarValue::UInt8(Some(v)) => v.to_string(),
1092        ScalarValue::UInt16(Some(v)) => v.to_string(),
1093        ScalarValue::UInt32(Some(v)) => v.to_string(),
1094        ScalarValue::UInt64(Some(v)) => v.to_string(),
1095        ScalarValue::Float32(Some(v)) => v.to_string(),
1096        ScalarValue::Float64(Some(v)) => v.to_string(),
1097        ScalarValue::Binary(Some(v))
1098        | ScalarValue::LargeBinary(Some(v))
1099        | ScalarValue::BinaryView(Some(v)) => String::from_utf8_lossy(v).into_owned(),
1100        ScalarValue::Date32(Some(days)) => {
1101            let date = NaiveDate::from_num_days_from_ce_opt(*days + 719_163);
1102            date.map(|d| d.format("%Y-%m-%d").to_string())
1103                .unwrap_or_default()
1104        }
1105        ScalarValue::TimestampSecond(Some(v), _) => format_unix(*v, 1),
1106        ScalarValue::TimestampMillisecond(Some(v), _) => format_unix(*v, 1_000),
1107        ScalarValue::TimestampMicrosecond(Some(v), _) => format_unix(*v, 1_000_000),
1108        ScalarValue::TimestampNanosecond(Some(v), _) => format_unix(*v, 1_000_000_000),
1109        _ => String::new(),
1110    }
1111}
1112
1113fn scalar_i64(v: &ScalarValue) -> Option<i64> {
1114    match v {
1115        ScalarValue::Int8(Some(v)) => Some(*v as i64),
1116        ScalarValue::Int16(Some(v)) => Some(*v as i64),
1117        ScalarValue::Int32(Some(v)) => Some(*v as i64),
1118        ScalarValue::Int64(Some(v)) => Some(*v),
1119        ScalarValue::UInt8(Some(v)) => Some(*v as i64),
1120        ScalarValue::UInt16(Some(v)) => Some(*v as i64),
1121        ScalarValue::UInt32(Some(v)) => Some(*v as i64),
1122        ScalarValue::UInt64(Some(v)) => i64::try_from(*v).ok(),
1123        ScalarValue::Float32(Some(v)) => Some(*v as i64),
1124        ScalarValue::Float64(Some(v)) => Some(*v as i64),
1125        ScalarValue::Utf8(Some(s))
1126        | ScalarValue::LargeUtf8(Some(s))
1127        | ScalarValue::Utf8View(Some(s)) => s.trim().parse().ok(),
1128        _ => None,
1129    }
1130}
1131
1132fn scalar_f64(v: &ScalarValue) -> Option<f64> {
1133    match v {
1134        ScalarValue::Float32(Some(v)) => Some(*v as f64),
1135        ScalarValue::Float64(Some(v)) => Some(*v),
1136        ScalarValue::Int8(Some(v)) => Some(*v as f64),
1137        ScalarValue::Int16(Some(v)) => Some(*v as f64),
1138        ScalarValue::Int32(Some(v)) => Some(*v as f64),
1139        ScalarValue::Int64(Some(v)) => Some(*v as f64),
1140        ScalarValue::UInt8(Some(v)) => Some(*v as f64),
1141        ScalarValue::UInt16(Some(v)) => Some(*v as f64),
1142        ScalarValue::UInt32(Some(v)) => Some(*v as f64),
1143        ScalarValue::UInt64(Some(v)) => Some(*v as f64),
1144        ScalarValue::Utf8(Some(s))
1145        | ScalarValue::LargeUtf8(Some(s))
1146        | ScalarValue::Utf8View(Some(s)) => s.trim().parse().ok(),
1147        _ => None,
1148    }
1149}
1150
1151fn scalar_bytes(v: &ScalarValue) -> Vec<u8> {
1152    match v {
1153        ScalarValue::Binary(Some(v))
1154        | ScalarValue::LargeBinary(Some(v))
1155        | ScalarValue::BinaryView(Some(v)) => v.clone(),
1156        _ => scalar_text(v).into_bytes(),
1157    }
1158}
1159
1160fn scalar_to_json_arg(v: ScalarValue) -> JsonValue {
1161    match v {
1162        ScalarValue::Boolean(Some(v)) => JsonValue::Bool(v),
1163        ScalarValue::Int8(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1164        ScalarValue::Int16(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1165        ScalarValue::Int32(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1166        ScalarValue::Int64(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1167        ScalarValue::UInt8(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1168        ScalarValue::UInt16(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1169        ScalarValue::UInt32(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1170        ScalarValue::UInt64(Some(v)) => JsonValue::Number(JsonNumber::from(v)),
1171        ScalarValue::Float32(Some(v)) => JsonNumber::from_f64(v as f64)
1172            .map(JsonValue::Number)
1173            .unwrap_or(JsonValue::Null),
1174        ScalarValue::Float64(Some(v)) => JsonNumber::from_f64(v)
1175            .map(JsonValue::Number)
1176            .unwrap_or(JsonValue::Null),
1177        ScalarValue::Utf8(Some(s))
1178        | ScalarValue::LargeUtf8(Some(s))
1179        | ScalarValue::Utf8View(Some(s)) => JsonValue::String(s),
1180        ScalarValue::Null => JsonValue::Null,
1181        _ if v.is_null() => JsonValue::Null,
1182        _ => JsonValue::String(scalar_text(&v)),
1183    }
1184}
1185
1186fn scalar_type_name(v: &ScalarValue) -> &'static str {
1187    if v.is_null() {
1188        return "null";
1189    }
1190    match v {
1191        ScalarValue::Boolean(_) => "integer",
1192        ScalarValue::Int8(_)
1193        | ScalarValue::Int16(_)
1194        | ScalarValue::Int32(_)
1195        | ScalarValue::Int64(_)
1196        | ScalarValue::UInt8(_)
1197        | ScalarValue::UInt16(_)
1198        | ScalarValue::UInt32(_)
1199        | ScalarValue::UInt64(_) => "integer",
1200        ScalarValue::Float16(_) | ScalarValue::Float32(_) | ScalarValue::Float64(_) => "real",
1201        ScalarValue::Binary(_) | ScalarValue::LargeBinary(_) | ScalarValue::BinaryView(_) => "blob",
1202        _ => "text",
1203    }
1204}
1205
1206fn abs_scalar(row: &RowArgs) -> ScalarValue {
1207    let Some(value) = row.scalar_any(0) else {
1208        return ScalarValue::Null;
1209    };
1210    match value {
1211        ScalarValue::Int8(Some(v)) => ScalarValue::Int8(v.checked_abs()),
1212        ScalarValue::Int16(Some(v)) => ScalarValue::Int16(v.checked_abs()),
1213        ScalarValue::Int32(Some(v)) => ScalarValue::Int32(v.checked_abs()),
1214        ScalarValue::Int64(Some(v)) => ScalarValue::Int64(v.checked_abs()),
1215        ScalarValue::UInt8(_)
1216        | ScalarValue::UInt16(_)
1217        | ScalarValue::UInt32(_)
1218        | ScalarValue::UInt64(_) => value,
1219        ScalarValue::Float32(Some(v)) => ScalarValue::Float32(Some(v.abs())),
1220        ScalarValue::Float64(Some(v)) => ScalarValue::Float64(Some(v.abs())),
1221        v if v.is_null() => v,
1222        v => scalar_f64(&v)
1223            .map(|n| ScalarValue::Float64(Some(n.abs())))
1224            .unwrap_or(ScalarValue::Float64(Some(0.0))),
1225    }
1226}
1227
1228fn coalesce_scalar(row: &RowArgs) -> ScalarValue {
1229    for idx in 0..row.len() {
1230        let value = row.scalar_any(idx).unwrap_or(ScalarValue::Null);
1231        if !value.is_null() {
1232            return value;
1233        }
1234    }
1235    ScalarValue::Null
1236}
1237
1238fn concat_row(row: &RowArgs) -> Option<String> {
1239    let mut out = String::new();
1240    for idx in 0..row.len() {
1241        if let Some(value) = row.scalar(idx) {
1242            out.push_str(&scalar_text(&value));
1243        }
1244    }
1245    Some(out)
1246}
1247
1248fn concat_ws_row(row: &RowArgs) -> Option<String> {
1249    let separator = row.text(0)?;
1250    let parts = (1..row.len())
1251        .filter_map(|idx| row.scalar(idx).map(|value| scalar_text(&value)))
1252        .collect::<Vec<_>>();
1253    Some(parts.join(&separator))
1254}
1255
1256fn scalar_extreme(row: &RowArgs, max: bool) -> ScalarValue {
1257    let mut best: Option<ScalarValue> = None;
1258    for idx in 0..row.len() {
1259        let value = row.scalar_any(idx).unwrap_or(ScalarValue::Null);
1260        if value.is_null() {
1261            return ScalarValue::Null;
1262        }
1263        let replace = best
1264            .as_ref()
1265            .map(|current| {
1266                let ordering = compare_scalar_values(&value, current);
1267                if max {
1268                    ordering == std::cmp::Ordering::Greater
1269                } else {
1270                    ordering == std::cmp::Ordering::Less
1271                }
1272            })
1273            .unwrap_or(true);
1274        if replace {
1275            best = Some(value);
1276        }
1277    }
1278    best.unwrap_or(ScalarValue::Null)
1279}
1280
1281fn compare_scalar_values(left: &ScalarValue, right: &ScalarValue) -> std::cmp::Ordering {
1282    match (scalar_f64(left), scalar_f64(right)) {
1283        (Some(a), Some(b)) => a.partial_cmp(&b).unwrap_or(std::cmp::Ordering::Equal),
1284        _ => scalar_text(left).cmp(&scalar_text(right)),
1285    }
1286}
1287
1288fn like_match_sql(value: &str, pattern: &str, escape: Option<char>) -> bool {
1289    fn rec(value: &[char], pattern: &[char], escape: Option<char>) -> bool {
1290        let Some((&head, rest)) = pattern.split_first() else {
1291            return value.is_empty();
1292        };
1293        if Some(head) == escape {
1294            let Some((&literal, rest)) = rest.split_first() else {
1295                return value.is_empty();
1296            };
1297            return value.first() == Some(&literal) && rec(&value[1..], rest, escape);
1298        }
1299        match head {
1300            '%' => {
1301                rec(value, rest, escape) || (!value.is_empty() && rec(&value[1..], pattern, escape))
1302            }
1303            '_' => !value.is_empty() && rec(&value[1..], rest, escape),
1304            ch => {
1305                value
1306                    .first()
1307                    .is_some_and(|value_ch| value_ch.eq_ignore_ascii_case(&ch))
1308                    && rec(&value[1..], rest, escape)
1309            }
1310        }
1311    }
1312    rec(
1313        &value.chars().collect::<Vec<_>>(),
1314        &pattern.chars().collect::<Vec<_>>(),
1315        escape,
1316    )
1317}
1318
1319fn compile_options() -> &'static [&'static str] {
1320    &[
1321        "DATAFUSION_54",
1322        "EXTENDED_SQL_FUNCTIONS",
1323        "JSON_TABLE_FUNCTIONS",
1324        "LOG_STRUCTURED_STORAGE",
1325        "MVCC",
1326        "WAL",
1327    ]
1328}
1329
1330fn math_unary(row: &RowArgs, f: impl FnOnce(f64) -> f64) -> Option<f64> {
1331    let out = f(row.f64(0)?);
1332    if out.is_nan() {
1333        None
1334    } else {
1335        Some(out)
1336    }
1337}
1338
1339fn log_row(row: &RowArgs) -> Option<f64> {
1340    let out = if row.len() == 1 {
1341        row.f64(0)?.log10()
1342    } else {
1343        row.f64(1)?.log(row.f64(0)?)
1344    };
1345    if out.is_nan() {
1346        None
1347    } else {
1348        Some(out)
1349    }
1350}
1351
1352fn soundex(input: &str) -> String {
1353    let mut letters = input.chars().filter(|ch| ch.is_ascii_alphabetic());
1354    let Some(first) = letters.next() else {
1355        return "?000".to_string();
1356    };
1357    let first = first.to_ascii_uppercase();
1358    let mut out = String::with_capacity(4);
1359    out.push(first);
1360    let mut last = soundex_code(first);
1361    for ch in letters {
1362        let code = soundex_code(ch.to_ascii_uppercase());
1363        if code != '0' && code != last {
1364            out.push(code);
1365            if out.len() == 4 {
1366                break;
1367            }
1368        }
1369        last = code;
1370    }
1371    while out.len() < 4 {
1372        out.push('0');
1373    }
1374    out
1375}
1376
1377fn soundex_code(ch: char) -> char {
1378    match ch {
1379        'B' | 'F' | 'P' | 'V' => '1',
1380        'C' | 'G' | 'J' | 'K' | 'Q' | 'S' | 'X' | 'Z' => '2',
1381        'D' | 'T' => '3',
1382        'L' => '4',
1383        'M' | 'N' => '5',
1384        'R' => '6',
1385        _ => '0',
1386    }
1387}
1388
1389fn unistr(input: &str) -> Option<String> {
1390    let chars = input.chars().collect::<Vec<_>>();
1391    let mut out = String::new();
1392    let mut idx = 0;
1393    while idx < chars.len() {
1394        let ch = chars[idx];
1395        idx += 1;
1396        if ch != '\\' {
1397            out.push(ch);
1398            continue;
1399        }
1400        let Some(marker) = chars.get(idx).copied() else {
1401            out.push('\\');
1402            break;
1403        };
1404        match marker {
1405            '\\' => {
1406                out.push('\\');
1407                idx += 1;
1408            }
1409            'u' => {
1410                idx += 1;
1411                out.push(read_hex_char(&chars, &mut idx, 4)?);
1412            }
1413            'U' => {
1414                idx += 1;
1415                out.push(read_hex_char(&chars, &mut idx, 8)?);
1416            }
1417            '+' => {
1418                idx += 1;
1419                out.push(read_hex_char(&chars, &mut idx, 6)?);
1420            }
1421            '0'..='9' | 'A'..='F' | 'a'..='f' => {
1422                out.push(read_hex_char(&chars, &mut idx, 4)?);
1423            }
1424            other => {
1425                out.push(other);
1426                idx += 1;
1427            }
1428        }
1429    }
1430    Some(out)
1431}
1432
1433fn read_hex_char(chars: &[char], idx: &mut usize, width: usize) -> Option<char> {
1434    if *idx + width > chars.len() {
1435        return None;
1436    }
1437    let mut value = 0_u32;
1438    for _ in 0..width {
1439        value = value.checked_mul(16)?;
1440        value = value.checked_add(chars[*idx].to_digit(16)?)?;
1441        *idx += 1;
1442    }
1443    char::from_u32(value)
1444}
1445
1446fn unistr_quote_text(input: &str) -> String {
1447    let mut out = String::new();
1448    for ch in input.chars() {
1449        match ch {
1450            '\u{0001}'..='\u{001f}' | '\\' => out.push_str(&format!("\\u{:04X}", ch as u32)),
1451            _ => out.push(ch),
1452        }
1453    }
1454    out
1455}
1456
1457fn format_unix(value: i64, scale: i64) -> String {
1458    let secs = value / scale;
1459    Utc.timestamp_opt(secs, 0)
1460        .single()
1461        .map(|dt| dt.format("%Y-%m-%d %H:%M:%S").to_string())
1462        .unwrap_or_default()
1463}
1464
1465fn parse_datetime_arg(row: &RowArgs, idx: usize) -> Option<DateTime<Utc>> {
1466    let mut dt = parse_datetime_at(row, idx)?;
1467    for i in (idx + 1)..row.len() {
1468        if let Some(modifier) = row.text(i) {
1469            dt = apply_datetime_modifier(dt, modifier.trim())?;
1470        }
1471    }
1472    Some(dt)
1473}
1474
1475fn parse_datetime_at(row: &RowArgs, idx: usize) -> Option<DateTime<Utc>> {
1476    match row.scalar(idx) {
1477        None => Some(Utc::now()),
1478        Some(v) => parse_datetime_scalar(&v),
1479    }
1480}
1481
1482fn parse_datetime_scalar(v: &ScalarValue) -> Option<DateTime<Utc>> {
1483    if let Some(seconds) = scalar_f64(v) {
1484        return unix_seconds_to_datetime(seconds);
1485    }
1486    let s = scalar_text(v);
1487    parse_datetime_text(s.trim())
1488}
1489
1490fn parse_datetime_text(s: &str) -> Option<DateTime<Utc>> {
1491    if s.is_empty() || s.eq_ignore_ascii_case("now") {
1492        return Some(Utc::now());
1493    }
1494    if let Ok(dt) = DateTime::parse_from_rfc3339(s) {
1495        return Some(dt.with_timezone(&Utc));
1496    }
1497    for fmt in ["%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S"] {
1498        if let Ok(dt) = NaiveDateTime::parse_from_str(s, fmt) {
1499            return Some(Utc.from_utc_datetime(&dt));
1500        }
1501    }
1502    if let Ok(date) = NaiveDate::parse_from_str(s, "%Y-%m-%d") {
1503        return date
1504            .and_hms_opt(0, 0, 0)
1505            .map(|dt| Utc.from_utc_datetime(&dt));
1506    }
1507    if let Ok(time) = NaiveTime::parse_from_str(s, "%H:%M:%S") {
1508        let date = NaiveDate::from_ymd_opt(1970, 1, 1)?;
1509        return Some(Utc.from_utc_datetime(&NaiveDateTime::new(date, time)));
1510    }
1511    s.parse::<f64>().ok().and_then(unix_seconds_to_datetime)
1512}
1513
1514fn unix_seconds_to_datetime(seconds: f64) -> Option<DateTime<Utc>> {
1515    if !seconds.is_finite() {
1516        return None;
1517    }
1518    let secs = seconds.trunc() as i64;
1519    let nanos = ((seconds.fract().abs()) * 1_000_000_000.0) as u32;
1520    Utc.timestamp_opt(secs, nanos).single()
1521}
1522
1523fn apply_datetime_modifier(dt: DateTime<Utc>, modifier: &str) -> Option<DateTime<Utc>> {
1524    if modifier.is_empty() || modifier.eq_ignore_ascii_case("unixepoch") {
1525        return Some(dt);
1526    }
1527    if modifier.eq_ignore_ascii_case("start of day") {
1528        return dt
1529            .date_naive()
1530            .and_hms_opt(0, 0, 0)
1531            .map(|v| Utc.from_utc_datetime(&v));
1532    }
1533    let mut parts = modifier.split_whitespace();
1534    let amount = parts.next()?.parse::<f64>().ok()?;
1535    let unit = parts.next()?.trim_end_matches('s').to_ascii_lowercase();
1536    let millis = match unit.as_str() {
1537        "day" => amount * 86_400_000.0,
1538        "hour" => amount * 3_600_000.0,
1539        "minute" => amount * 60_000.0,
1540        "second" => amount * 1_000.0,
1541        _ => return None,
1542    };
1543    Some(dt + Duration::milliseconds(millis as i64))
1544}
1545
1546fn unix_to_julian(seconds: f64) -> f64 {
1547    seconds / 86_400.0 + 2_440_587.5
1548}
1549
1550fn strftime_row(row: &RowArgs) -> Option<String> {
1551    let fmt = row.text(0)?;
1552    let dt = parse_datetime_arg(row, 1)?;
1553    let mut out = String::new();
1554    let mut chars = fmt.chars();
1555    while let Some(ch) = chars.next() {
1556        if ch != '%' {
1557            out.push(ch);
1558            continue;
1559        }
1560        match chars.next()? {
1561            '%' => out.push('%'),
1562            'Y' => out.push_str(&format!("{:04}", dt.year())),
1563            'm' => out.push_str(&format!("{:02}", dt.month())),
1564            'd' => out.push_str(&format!("{:02}", dt.day())),
1565            'H' => out.push_str(&format!("{:02}", dt.hour())),
1566            'M' => out.push_str(&format!("{:02}", dt.minute())),
1567            'S' => out.push_str(&format!("{:02}", dt.second())),
1568            'f' => out.push_str(&format!(
1569                "{:02}.{:03}",
1570                dt.second(),
1571                dt.timestamp_subsec_millis()
1572            )),
1573            's' => out.push_str(&dt.timestamp().to_string()),
1574            'J' => out.push_str(&unix_to_julian(dt.timestamp() as f64).to_string()),
1575            other => {
1576                out.push('%');
1577                out.push(other);
1578            }
1579        }
1580    }
1581    Some(out)
1582}
1583
1584fn parse_json_arg(row: &RowArgs, idx: usize) -> Option<JsonValue> {
1585    serde_json::from_str(&row.text(idx)?).ok()
1586}
1587
1588fn json_extract_row(row: &RowArgs) -> Option<String> {
1589    let root = parse_json_arg(row, 0)?;
1590    if row.len() <= 1 {
1591        return json_scalar_result(&root);
1592    }
1593    let mut values = Vec::new();
1594    for idx in 1..row.len() {
1595        let path = row.text(idx)?;
1596        let found = json_path_get(&root, &parse_json_path(&path)?)
1597            .cloned()
1598            .unwrap_or(JsonValue::Null);
1599        if row.len() == 2 {
1600            return json_scalar_result(&found);
1601        }
1602        values.push(found);
1603    }
1604    Some(JsonValue::Array(values).to_string())
1605}
1606
1607fn json_type_row(row: &RowArgs) -> Option<String> {
1608    let root = parse_json_arg(row, 0)?;
1609    let value = if row.len() > 1 {
1610        let path = row.text(1)?;
1611        json_path_get(&root, &parse_json_path(&path)?)?
1612    } else {
1613        &root
1614    };
1615    Some(json_type(value).to_string())
1616}
1617
1618fn json_array_length_row(row: &RowArgs) -> Option<i64> {
1619    let root = parse_json_arg(row, 0)?;
1620    let value = if row.len() > 1 {
1621        let path = row.text(1)?;
1622        json_path_get(&root, &parse_json_path(&path)?)?
1623    } else {
1624        &root
1625    };
1626    Some(value.as_array().map(|v| v.len() as i64).unwrap_or(0))
1627}
1628
1629#[derive(Clone, Copy)]
1630enum JsonMutation {
1631    Set,
1632    Insert,
1633    Replace,
1634}
1635
1636fn json_mutate_row(row: &RowArgs, mode: JsonMutation) -> Option<String> {
1637    if row.len() < 3 {
1638        return None;
1639    }
1640    let mut root = parse_json_arg(row, 0)?;
1641    let mut idx = 1;
1642    while idx + 1 < row.len() {
1643        let path = parse_json_path(&row.text(idx)?)?;
1644        let value = row.scalar(idx + 1).map(scalar_to_json_arg)?;
1645        match mode {
1646            JsonMutation::Set => json_path_set(&mut root, &path, value, true),
1647            JsonMutation::Insert => json_path_set(&mut root, &path, value, false),
1648            JsonMutation::Replace => {
1649                if json_path_get(&root, &path).is_some() {
1650                    json_path_set(&mut root, &path, value, true);
1651                }
1652            }
1653        }
1654        idx += 2;
1655    }
1656    Some(root.to_string())
1657}
1658
1659fn json_array_insert_row(row: &RowArgs) -> Option<String> {
1660    if row.len() < 3 {
1661        return None;
1662    }
1663    let mut root = parse_json_arg(row, 0)?;
1664    let mut idx = 1;
1665    while idx + 1 < row.len() {
1666        let path = parse_json_path(&row.text(idx)?)?;
1667        let value = row.scalar(idx + 1).map(scalar_to_json_arg)?;
1668        json_path_array_insert(&mut root, &path, value);
1669        idx += 2;
1670    }
1671    Some(root.to_string())
1672}
1673
1674fn json_path_array_insert(value: &mut JsonValue, path: &[JsonPathToken], new_value: JsonValue) {
1675    let Some((last, parent_path)) = path.split_last() else {
1676        return;
1677    };
1678    let mut cur = value;
1679    for token in parent_path {
1680        cur = match token {
1681            JsonPathToken::Key(k) => match cur.as_object_mut().and_then(|obj| obj.get_mut(k)) {
1682                Some(value) => value,
1683                None => return,
1684            },
1685            JsonPathToken::Index(i) => match cur.as_array_mut().and_then(|arr| arr.get_mut(*i)) {
1686                Some(value) => value,
1687                None => return,
1688            },
1689        };
1690    }
1691    let JsonPathToken::Index(index) = last else {
1692        return;
1693    };
1694    if let Some(arr) = cur.as_array_mut() {
1695        if *index <= arr.len() {
1696            arr.insert(*index, new_value);
1697        }
1698    }
1699}
1700
1701fn json_array_row(row: &RowArgs) -> Option<String> {
1702    let values = (0..row.len())
1703        .map(|idx| scalar_to_json_arg(row.scalar_any(idx).unwrap_or(ScalarValue::Null)))
1704        .collect::<Vec<_>>();
1705    Some(JsonValue::Array(values).to_string())
1706}
1707
1708fn json_object_row(row: &RowArgs) -> Option<String> {
1709    if row.len() % 2 != 0 {
1710        return None;
1711    }
1712    let mut obj = JsonMap::new();
1713    let mut idx = 0;
1714    while idx + 1 < row.len() {
1715        let key = row.text(idx)?;
1716        let value = scalar_to_json_arg(row.scalar_any(idx + 1).unwrap_or(ScalarValue::Null));
1717        obj.insert(key, value);
1718        idx += 2;
1719    }
1720    Some(JsonValue::Object(obj).to_string())
1721}
1722
1723fn json_patch_row(row: &RowArgs) -> Option<String> {
1724    let target = parse_json_arg(row, 0)?;
1725    let patch = parse_json_arg(row, 1)?;
1726    Some(json_merge_patch(target, patch).to_string())
1727}
1728
1729fn json_merge_patch(target: JsonValue, patch: JsonValue) -> JsonValue {
1730    let JsonValue::Object(patch_obj) = patch else {
1731        return patch;
1732    };
1733    let mut target_obj = match target {
1734        JsonValue::Object(obj) => obj,
1735        _ => JsonMap::new(),
1736    };
1737    for (key, value) in patch_obj {
1738        if value.is_null() {
1739            target_obj.remove(&key);
1740        } else {
1741            let current = target_obj.remove(&key).unwrap_or(JsonValue::Null);
1742            target_obj.insert(key, json_merge_patch(current, value));
1743        }
1744    }
1745    JsonValue::Object(target_obj)
1746}
1747
1748fn json_error_position(input: &str) -> usize {
1749    match serde_json::from_str::<JsonValue>(input) {
1750        Ok(_) => 0,
1751        Err(err) => {
1752            let line = err.line().max(1);
1753            let column = err.column();
1754            let mut current_line = 1;
1755            let mut line_start = 0;
1756            for (idx, ch) in input.char_indices() {
1757                if current_line == line {
1758                    break;
1759                }
1760                if ch == '\n' {
1761                    current_line += 1;
1762                    line_start = idx + ch.len_utf8();
1763                }
1764            }
1765            line_start + column.max(1)
1766        }
1767    }
1768}
1769
1770fn json_remove_row(row: &RowArgs) -> Option<String> {
1771    let mut root = parse_json_arg(row, 0)?;
1772    for idx in 1..row.len() {
1773        let path = parse_json_path(&row.text(idx)?)?;
1774        json_path_remove(&mut root, &path);
1775    }
1776    Some(root.to_string())
1777}
1778
1779fn json_scalar_result(value: &JsonValue) -> Option<String> {
1780    match value {
1781        JsonValue::Null => None,
1782        JsonValue::Bool(true) => Some("1".into()),
1783        JsonValue::Bool(false) => Some("0".into()),
1784        JsonValue::Number(n) => Some(n.to_string()),
1785        JsonValue::String(s) => Some(s.clone()),
1786        JsonValue::Array(_) | JsonValue::Object(_) => Some(value.to_string()),
1787    }
1788}
1789
1790fn json_type(value: &JsonValue) -> &'static str {
1791    match value {
1792        JsonValue::Null => "null",
1793        JsonValue::Bool(true) => "true",
1794        JsonValue::Bool(false) => "false",
1795        JsonValue::Number(n) if n.is_i64() || n.is_u64() => "integer",
1796        JsonValue::Number(_) => "real",
1797        JsonValue::String(_) => "text",
1798        JsonValue::Array(_) => "array",
1799        JsonValue::Object(_) => "object",
1800    }
1801}
1802
1803#[derive(Debug, Clone, PartialEq, Eq)]
1804enum JsonPathToken {
1805    Key(String),
1806    Index(usize),
1807}
1808
1809fn parse_json_path(path: &str) -> Option<Vec<JsonPathToken>> {
1810    let chars = path.chars().collect::<Vec<_>>();
1811    if chars.first() != Some(&'$') {
1812        return None;
1813    }
1814    let mut idx = 1;
1815    let mut tokens = Vec::new();
1816    while idx < chars.len() {
1817        match chars[idx] {
1818            '.' => {
1819                idx += 1;
1820                let mut key = String::new();
1821                while idx < chars.len() {
1822                    let c = chars[idx];
1823                    if c == '.' || c == '[' {
1824                        break;
1825                    }
1826                    key.push(c);
1827                    idx += 1;
1828                }
1829                if key.is_empty() {
1830                    return None;
1831                }
1832                tokens.push(JsonPathToken::Key(key));
1833            }
1834            '[' => {
1835                idx += 1;
1836                if idx >= chars.len() {
1837                    return None;
1838                }
1839                if matches!(chars[idx], '"' | '\'') {
1840                    let quote = chars[idx];
1841                    idx += 1;
1842                    let mut key = String::new();
1843                    while idx < chars.len() {
1844                        let c = chars[idx];
1845                        idx += 1;
1846                        if c == quote {
1847                            break;
1848                        }
1849                        if c == '\\' && idx < chars.len() {
1850                            key.push(chars[idx]);
1851                            idx += 1;
1852                        } else {
1853                            key.push(c);
1854                        }
1855                    }
1856                    if idx >= chars.len() || chars[idx] != ']' {
1857                        return None;
1858                    }
1859                    idx += 1;
1860                    tokens.push(JsonPathToken::Key(key));
1861                } else {
1862                    let mut n = String::new();
1863                    while idx < chars.len() && chars[idx] != ']' {
1864                        n.push(chars[idx]);
1865                        idx += 1;
1866                    }
1867                    if idx >= chars.len() || chars[idx] != ']' {
1868                        return None;
1869                    }
1870                    idx += 1;
1871                    tokens.push(JsonPathToken::Index(n.parse().ok()?));
1872                }
1873            }
1874            _ => return None,
1875        }
1876    }
1877    Some(tokens)
1878}
1879
1880fn json_path_get<'a>(value: &'a JsonValue, path: &[JsonPathToken]) -> Option<&'a JsonValue> {
1881    let mut cur = value;
1882    for token in path {
1883        cur = match token {
1884            JsonPathToken::Key(k) => cur.as_object()?.get(k)?,
1885            JsonPathToken::Index(i) => cur.as_array()?.get(*i)?,
1886        };
1887    }
1888    Some(cur)
1889}
1890
1891fn json_path_set(
1892    value: &mut JsonValue,
1893    path: &[JsonPathToken],
1894    new_value: JsonValue,
1895    overwrite: bool,
1896) {
1897    if path.is_empty() {
1898        if overwrite {
1899            *value = new_value;
1900        }
1901        return;
1902    }
1903    let mut cur = value;
1904    for token in &path[..path.len() - 1] {
1905        match token {
1906            JsonPathToken::Key(k) => {
1907                if !cur.is_object() {
1908                    *cur = JsonValue::Object(JsonMap::new());
1909                }
1910                cur = cur
1911                    .as_object_mut()
1912                    .unwrap()
1913                    .entry(k.clone())
1914                    .or_insert(JsonValue::Object(JsonMap::new()));
1915            }
1916            JsonPathToken::Index(i) => {
1917                if let Some(arr) = cur.as_array_mut() {
1918                    if *i >= arr.len() {
1919                        return;
1920                    }
1921                    cur = &mut arr[*i];
1922                } else {
1923                    return;
1924                }
1925            }
1926        }
1927    }
1928    match path.last().unwrap() {
1929        JsonPathToken::Key(k) => {
1930            if !cur.is_object() {
1931                *cur = JsonValue::Object(JsonMap::new());
1932            }
1933            let obj = cur.as_object_mut().unwrap();
1934            if overwrite || !obj.contains_key(k) {
1935                obj.insert(k.clone(), new_value);
1936            }
1937        }
1938        JsonPathToken::Index(i) => {
1939            if let Some(arr) = cur.as_array_mut() {
1940                if *i < arr.len() {
1941                    if overwrite {
1942                        arr[*i] = new_value;
1943                    }
1944                } else if *i == arr.len() {
1945                    arr.push(new_value);
1946                }
1947            }
1948        }
1949    }
1950}
1951
1952fn json_path_remove(value: &mut JsonValue, path: &[JsonPathToken]) {
1953    if path.is_empty() {
1954        return;
1955    }
1956    let mut cur = value;
1957    for token in &path[..path.len() - 1] {
1958        cur = match token {
1959            JsonPathToken::Key(k) => match cur.as_object_mut().and_then(|o| o.get_mut(k)) {
1960                Some(v) => v,
1961                None => return,
1962            },
1963            JsonPathToken::Index(i) => match cur.as_array_mut().and_then(|a| a.get_mut(*i)) {
1964                Some(v) => v,
1965                None => return,
1966            },
1967        };
1968    }
1969    match path.last().unwrap() {
1970        JsonPathToken::Key(k) => {
1971            if let Some(obj) = cur.as_object_mut() {
1972                obj.remove(k);
1973            }
1974        }
1975        JsonPathToken::Index(i) => {
1976            if let Some(arr) = cur.as_array_mut() {
1977                if *i < arr.len() {
1978                    arr.remove(*i);
1979                }
1980            }
1981        }
1982    }
1983}
1984
1985fn instr(haystack: &str, needle: &str) -> i64 {
1986    if needle.is_empty() {
1987        return 1;
1988    }
1989    haystack
1990        .find(needle)
1991        .map(|byte| haystack[..byte].chars().count() as i64 + 1)
1992        .unwrap_or(0)
1993}
1994
1995fn sql_quote(v: ScalarValue) -> String {
1996    if v.is_null() {
1997        return "NULL".into();
1998    }
1999    match v {
2000        ScalarValue::Utf8(Some(s))
2001        | ScalarValue::LargeUtf8(Some(s))
2002        | ScalarValue::Utf8View(Some(s)) => format!("'{}'", s.replace('\'', "''")),
2003        ScalarValue::Binary(Some(b))
2004        | ScalarValue::LargeBinary(Some(b))
2005        | ScalarValue::BinaryView(Some(b)) => format!("X'{}'", hex_encode(&b)),
2006        _ => scalar_text(&v),
2007    }
2008}
2009
2010fn hex_encode(bytes: &[u8]) -> String {
2011    const HEX: &[u8; 16] = b"0123456789ABCDEF";
2012    let mut out = String::with_capacity(bytes.len() * 2);
2013    for b in bytes {
2014        out.push(HEX[(b >> 4) as usize] as char);
2015        out.push(HEX[(b & 0x0f) as usize] as char);
2016    }
2017    out
2018}
2019
2020fn hex_decode_ignoring(value: &str, ignored: &str) -> Option<Vec<u8>> {
2021    let value = value.trim();
2022    let filtered = value
2023        .bytes()
2024        .filter(|b| !ignored.as_bytes().contains(b))
2025        .collect::<Vec<_>>();
2026    if filtered.len() % 2 != 0 {
2027        return None;
2028    }
2029    let mut out = Vec::with_capacity(filtered.len() / 2);
2030    for i in (0..filtered.len()).step_by(2) {
2031        let hi = hex_value(filtered[i])?;
2032        let lo = hex_value(filtered[i + 1])?;
2033        out.push((hi << 4) | lo);
2034    }
2035    Some(out)
2036}
2037
2038fn hex_value(b: u8) -> Option<u8> {
2039    match b {
2040        b'0'..=b'9' => Some(b - b'0'),
2041        b'a'..=b'f' => Some(b - b'a' + 10),
2042        b'A'..=b'F' => Some(b - b'A' + 10),
2043        _ => None,
2044    }
2045}
2046
2047fn printf_row(row: &RowArgs) -> Option<String> {
2048    let format = row.text(0)?;
2049    let mut out = String::new();
2050    let mut arg_idx = 1;
2051    let chars = format.chars().collect::<Vec<_>>();
2052    let mut idx = 0;
2053    while idx < chars.len() {
2054        let ch = chars[idx];
2055        idx += 1;
2056        if ch != '%' {
2057            out.push(ch);
2058            continue;
2059        }
2060        if idx >= chars.len() {
2061            out.push('%');
2062            break;
2063        }
2064        if chars[idx] == '%' {
2065            out.push('%');
2066            idx += 1;
2067            continue;
2068        }
2069        let spec_start = idx;
2070        let flags = PrintfFlags::parse(&chars, &mut idx);
2071        let width = parse_decimal(&chars, &mut idx);
2072        let precision = if idx < chars.len() && chars[idx] == '.' {
2073            idx += 1;
2074            Some(parse_decimal(&chars, &mut idx).unwrap_or(0))
2075        } else {
2076            None
2077        };
2078        while idx < chars.len() && matches!(chars[idx], 'l' | 'z' | 't' | 'j') {
2079            idx += 1;
2080        }
2081        if idx >= chars.len() {
2082            out.push('%');
2083            for ch in &chars[spec_start..] {
2084                out.push(*ch);
2085            }
2086            break;
2087        }
2088        let spec = chars[idx];
2089        idx += 1;
2090        let formatted = match spec {
2091            's' | 'z' | 'q' | 'Q' => {
2092                let value = row.scalar(arg_idx).unwrap_or(ScalarValue::Null);
2093                arg_idx += 1;
2094                Some(format_printf_text(value, spec, precision))
2095            }
2096            'd' | 'i' | 'u' => {
2097                let value = row.i64(arg_idx).unwrap_or(0);
2098                arg_idx += 1;
2099                Some(format_printf_int(value, 10, false, &flags, precision))
2100            }
2101            'x' | 'X' => {
2102                let value = row.i64(arg_idx).unwrap_or(0);
2103                arg_idx += 1;
2104                Some(format_printf_int(value, 16, spec == 'X', &flags, precision))
2105            }
2106            'o' => {
2107                let value = row.i64(arg_idx).unwrap_or(0);
2108                arg_idx += 1;
2109                Some(format_printf_int(value, 8, false, &flags, precision))
2110            }
2111            'c' => {
2112                let value = row.scalar(arg_idx).unwrap_or(ScalarValue::Null);
2113                arg_idx += 1;
2114                Some(format_printf_char(value, precision.unwrap_or(1)))
2115            }
2116            'f' | 'F' | 'e' | 'E' | 'g' | 'G' => {
2117                let value = row.f64(arg_idx).unwrap_or(0.0);
2118                arg_idx += 1;
2119                Some(format_printf_float(value, spec, &flags, precision))
2120            }
2121            'n' => {
2122                arg_idx += 1;
2123                Some(String::new())
2124            }
2125            other => Some(format!("%{other}")),
2126        }
2127        .unwrap_or_default();
2128        out.push_str(&apply_printf_width(
2129            formatted,
2130            width,
2131            flags.left,
2132            flags.zero && !flags.left,
2133        ));
2134    }
2135    Some(out)
2136}
2137
2138#[derive(Default)]
2139struct PrintfFlags {
2140    left: bool,
2141    plus: bool,
2142    space: bool,
2143    alternate: bool,
2144    zero: bool,
2145}
2146
2147impl PrintfFlags {
2148    fn parse(chars: &[char], idx: &mut usize) -> Self {
2149        let mut flags = Self::default();
2150        while *idx < chars.len() {
2151            match chars[*idx] {
2152                '-' => flags.left = true,
2153                '+' => flags.plus = true,
2154                ' ' => flags.space = true,
2155                '#' => flags.alternate = true,
2156                '0' => flags.zero = true,
2157                ',' | '!' => {}
2158                _ => break,
2159            }
2160            *idx += 1;
2161        }
2162        flags
2163    }
2164}
2165
2166fn parse_decimal(chars: &[char], idx: &mut usize) -> Option<usize> {
2167    let start = *idx;
2168    let mut value = 0_usize;
2169    while *idx < chars.len() {
2170        let Some(digit) = chars[*idx].to_digit(10) else {
2171            break;
2172        };
2173        value = value.saturating_mul(10).saturating_add(digit as usize);
2174        *idx += 1;
2175    }
2176    if *idx == start {
2177        None
2178    } else {
2179        Some(value)
2180    }
2181}
2182
2183fn format_printf_text(value: ScalarValue, spec: char, precision: Option<usize>) -> String {
2184    if spec == 'Q' && value.is_null() {
2185        return "NULL".to_string();
2186    }
2187    let mut text = scalar_text(&value);
2188    if matches!(spec, 'q' | 'Q') {
2189        text = text.replace('\'', "''");
2190    }
2191    if let Some(limit) = precision {
2192        text = text.chars().take(limit).collect();
2193    }
2194    if spec == 'Q' {
2195        format!("'{text}'")
2196    } else {
2197        text
2198    }
2199}
2200
2201fn format_printf_int(
2202    value: i64,
2203    radix: u32,
2204    uppercase: bool,
2205    flags: &PrintfFlags,
2206    precision: Option<usize>,
2207) -> String {
2208    let negative = value < 0 && radix == 10;
2209    let magnitude = if negative {
2210        value.unsigned_abs()
2211    } else {
2212        value as u64
2213    };
2214    let mut digits = match radix {
2215        8 => format!("{magnitude:o}"),
2216        16 if uppercase => format!("{magnitude:X}"),
2217        16 => format!("{magnitude:x}"),
2218        _ => magnitude.to_string(),
2219    };
2220    if let Some(precision) = precision {
2221        if precision == 0 && magnitude == 0 {
2222            digits.clear();
2223        } else if digits.len() < precision {
2224            digits = format!("{}{}", "0".repeat(precision - digits.len()), digits);
2225        }
2226    }
2227    let mut out = String::new();
2228    if negative {
2229        out.push('-');
2230    } else if flags.plus && radix == 10 {
2231        out.push('+');
2232    } else if flags.space && radix == 10 {
2233        out.push(' ');
2234    }
2235    if flags.alternate {
2236        match radix {
2237            8 if !digits.starts_with('0') => out.push('0'),
2238            16 if uppercase && magnitude != 0 => out.push_str("0X"),
2239            16 if magnitude != 0 => out.push_str("0x"),
2240            _ => {}
2241        }
2242    }
2243    out.push_str(&digits);
2244    out
2245}
2246
2247fn format_printf_float(
2248    value: f64,
2249    spec: char,
2250    flags: &PrintfFlags,
2251    precision: Option<usize>,
2252) -> String {
2253    let precision = precision.unwrap_or(6);
2254    let mut out = match spec {
2255        'e' => format!("{value:.precision$e}"),
2256        'E' => format!("{value:.precision$E}"),
2257        'g' | 'G' => {
2258            let abs = value.abs();
2259            if abs != 0.0 && !(0.0001..1_000_000.0).contains(&abs) {
2260                if spec == 'G' {
2261                    format!("{value:.precision$E}")
2262                } else {
2263                    format!("{value:.precision$e}")
2264                }
2265            } else {
2266                let text = format!("{value:.precision$}");
2267                text.trim_end_matches('0').trim_end_matches('.').to_string()
2268            }
2269        }
2270        _ => format!("{value:.precision$}"),
2271    };
2272    if value >= 0.0 {
2273        if flags.plus {
2274            out.insert(0, '+');
2275        } else if flags.space {
2276            out.insert(0, ' ');
2277        }
2278    }
2279    if flags.alternate && !out.contains('.') && !out.contains('e') && !out.contains('E') {
2280        out.push('.');
2281    }
2282    out
2283}
2284
2285fn format_printf_char(value: ScalarValue, repeat: usize) -> String {
2286    let ch = match value {
2287        ScalarValue::Utf8(Some(s))
2288        | ScalarValue::LargeUtf8(Some(s))
2289        | ScalarValue::Utf8View(Some(s)) => s.chars().next().unwrap_or('\0'),
2290        other => scalar_i64(&other)
2291            .and_then(|v| char::from_u32(v as u32))
2292            .unwrap_or('\0'),
2293    };
2294    ch.to_string().repeat(repeat.max(1))
2295}
2296
2297fn apply_printf_width(mut text: String, width: Option<usize>, left: bool, zero: bool) -> String {
2298    let Some(width) = width else {
2299        return text;
2300    };
2301    let len = text.chars().count();
2302    if len >= width {
2303        return text;
2304    }
2305    let pad = width - len;
2306    let ch = if zero { '0' } else { ' ' };
2307    if left {
2308        text.extend(std::iter::repeat(ch).take(pad));
2309        text
2310    } else if zero && matches!(text.as_bytes().first(), Some(b'+' | b'-' | b' ')) {
2311        let sign = text.remove(0);
2312        format!("{sign}{}{}", ch.to_string().repeat(pad), text)
2313    } else if zero && (text.starts_with("0x") || text.starts_with("0X")) {
2314        let prefix = text[..2].to_string();
2315        let rest = text[2..].to_string();
2316        format!("{prefix}{}{}", ch.to_string().repeat(pad), rest)
2317    } else {
2318        format!("{}{}", ch.to_string().repeat(pad), text)
2319    }
2320}
2321
2322fn char_row(row: &RowArgs) -> Option<String> {
2323    let mut out = String::new();
2324    for idx in 0..row.len() {
2325        if let Some(code) = row.i64(idx) {
2326            if let Some(ch) = char::from_u32(code as u32) {
2327                out.push(ch);
2328            }
2329        }
2330    }
2331    Some(out)
2332}
2333
2334fn substr_row(row: &RowArgs) -> Option<String> {
2335    let value = row.text(0)?;
2336    let chars = value.chars().collect::<Vec<_>>();
2337    let start = row.i64(1)?;
2338    let len = row.i64(2);
2339    let start_idx = if start > 0 {
2340        (start - 1) as usize
2341    } else if start < 0 {
2342        chars.len().saturating_sub(start.unsigned_abs() as usize)
2343    } else {
2344        0
2345    };
2346    if start_idx >= chars.len() {
2347        return Some(String::new());
2348    }
2349    let end = len
2350        .map(|n| start_idx.saturating_add(n.max(0) as usize))
2351        .unwrap_or(chars.len())
2352        .min(chars.len());
2353    Some(chars[start_idx..end].iter().collect())
2354}
2355
2356fn glob_match(pattern: &str, value: &str) -> bool {
2357    fn rec(pattern: &[char], value: &[char]) -> bool {
2358        match pattern.split_first() {
2359            None => value.is_empty(),
2360            Some(('*', rest)) => {
2361                rec(rest, value) || (!value.is_empty() && rec(pattern, &value[1..]))
2362            }
2363            Some(('?', rest)) => !value.is_empty() && rec(rest, &value[1..]),
2364            Some(('[', _)) => {
2365                let Some((matched, rest)) = glob_class_match(pattern, value.first().copied())
2366                else {
2367                    return value.first() == Some(&'[') && rec(&pattern[1..], &value[1..]);
2368                };
2369                matched && !value.is_empty() && rec(rest, &value[1..])
2370            }
2371            Some((ch, rest)) => value.first() == Some(ch) && rec(rest, &value[1..]),
2372        }
2373    }
2374    rec(
2375        &pattern.chars().collect::<Vec<_>>(),
2376        &value.chars().collect::<Vec<_>>(),
2377    )
2378}
2379
2380fn glob_class_match(pattern: &[char], value: Option<char>) -> Option<(bool, &[char])> {
2381    if pattern.first() != Some(&'[') {
2382        return None;
2383    }
2384    let value = value?;
2385    let mut idx = 1;
2386    let mut negated = false;
2387    if matches!(pattern.get(idx), Some('^' | '!')) {
2388        negated = true;
2389        idx += 1;
2390    }
2391    let mut matched = false;
2392    let mut saw_entry = false;
2393    while idx < pattern.len() {
2394        if pattern[idx] == ']' && saw_entry {
2395            return Some((
2396                if negated { !matched } else { matched },
2397                &pattern[idx + 1..],
2398            ));
2399        }
2400        let start = pattern[idx];
2401        saw_entry = true;
2402        if idx + 2 < pattern.len() && pattern[idx + 1] == '-' && pattern[idx + 2] != ']' {
2403            let end = pattern[idx + 2];
2404            if start <= value && value <= end {
2405                matched = true;
2406            }
2407            idx += 3;
2408        } else {
2409            if value == start {
2410                matched = true;
2411            }
2412            idx += 1;
2413        }
2414    }
2415    None
2416}
2417
2418fn random_bytes(len: usize) -> Vec<u8> {
2419    let mut out = Vec::with_capacity(len);
2420    while out.len() < len {
2421        out.extend_from_slice(&next_random().to_ne_bytes());
2422    }
2423    out.truncate(len);
2424    out
2425}
2426
2427fn next_random() -> i64 {
2428    let now = std::time::SystemTime::now()
2429        .duration_since(std::time::UNIX_EPOCH)
2430        .map(|d| d.as_nanos() as u64)
2431        .unwrap_or(0);
2432    let old = RANDOM_STATE.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::Relaxed);
2433    let mut x = old ^ now;
2434    x ^= x >> 12;
2435    x ^= x << 25;
2436    x ^= x >> 27;
2437    let value = (x.wrapping_mul(0x2545_F491_4F6C_DD1D)) as i64;
2438    if value == i64::MIN {
2439        0
2440    } else {
2441        value
2442    }
2443}
2444
2445#[derive(Debug)]
2446struct JsonEachFunc;
2447
2448impl TableFunctionImpl for JsonEachFunc {
2449    fn call_with_args(&self, args: TableFunctionArgs) -> DFResult<Arc<dyn TableProvider>> {
2450        json_table_function(args, "json_each", JsonTableMode::Each)
2451    }
2452}
2453
2454#[derive(Debug)]
2455struct JsonTreeFunc;
2456
2457impl TableFunctionImpl for JsonTreeFunc {
2458    fn call_with_args(&self, args: TableFunctionArgs) -> DFResult<Arc<dyn TableProvider>> {
2459        json_table_function(args, "json_tree", JsonTableMode::Tree)
2460    }
2461}
2462
2463#[derive(Debug)]
2464struct JsonbEachFunc;
2465
2466impl TableFunctionImpl for JsonbEachFunc {
2467    fn call_with_args(&self, args: TableFunctionArgs) -> DFResult<Arc<dyn TableProvider>> {
2468        json_table_function(args, "jsonb_each", JsonTableMode::Each)
2469    }
2470}
2471
2472#[derive(Debug)]
2473struct JsonbTreeFunc;
2474
2475impl TableFunctionImpl for JsonbTreeFunc {
2476    fn call_with_args(&self, args: TableFunctionArgs) -> DFResult<Arc<dyn TableProvider>> {
2477        json_table_function(args, "jsonb_tree", JsonTableMode::Tree)
2478    }
2479}
2480
2481#[derive(Debug)]
2482struct SeriesFunc;
2483
2484impl TableFunctionImpl for SeriesFunc {
2485    fn call_with_args(&self, args: TableFunctionArgs) -> DFResult<Arc<dyn TableProvider>> {
2486        let exprs = args.exprs();
2487        let mut values = Vec::with_capacity(exprs.len());
2488        for expr in exprs {
2489            values.push(literal_i64(expr).ok_or_else(|| {
2490                DataFusionError::Plan("series arguments must be integer literals".into())
2491            })?);
2492        }
2493        let (start, stop, step) = match values.as_slice() {
2494            [stop] => (0, *stop, 1),
2495            [start, stop] => (*start, *stop, 1),
2496            [start, stop, step] => (*start, *stop, *step),
2497            _ => {
2498                return Err(DataFusionError::Plan(
2499                    "series requires stop, start/stop, or start/stop/step".into(),
2500                ));
2501            }
2502        };
2503        series_provider(start, stop, step)
2504    }
2505}
2506
2507fn literal_string(expr: &DfExpr) -> Option<String> {
2508    match expr {
2509        DfExpr::Literal(
2510            ScalarValue::Utf8(Some(s))
2511            | ScalarValue::LargeUtf8(Some(s))
2512            | ScalarValue::Utf8View(Some(s)),
2513            _,
2514        ) => Some(s.clone()),
2515        _ => None,
2516    }
2517}
2518
2519fn literal_i64(expr: &DfExpr) -> Option<i64> {
2520    match expr {
2521        DfExpr::Literal(ScalarValue::Int64(Some(v)), _) => Some(*v),
2522        DfExpr::Literal(ScalarValue::Int32(Some(v)), _) => Some(*v as i64),
2523        DfExpr::Literal(ScalarValue::UInt64(Some(v)), _) => i64::try_from(*v).ok(),
2524        DfExpr::Literal(ScalarValue::UInt32(Some(v)), _) => Some(*v as i64),
2525        _ => None,
2526    }
2527}
2528
2529pub(crate) fn series_provider(
2530    start: i64,
2531    stop: i64,
2532    step: i64,
2533) -> DFResult<Arc<dyn TableProvider>> {
2534    if step == 0 {
2535        return Err(DataFusionError::Plan("series step must not be 0".into()));
2536    }
2537    let mut values = Vec::new();
2538    let mut current = start;
2539    while if step > 0 {
2540        current <= stop
2541    } else {
2542        current >= stop
2543    } {
2544        if values.len() >= 1_000_000 {
2545            return Err(DataFusionError::Plan(
2546                "series output is capped at 1,000,000 rows".into(),
2547            ));
2548        }
2549        values.push(current);
2550        current = current.saturating_add(step);
2551        if (step > 0 && current == i64::MAX) || (step < 0 && current == i64::MIN) {
2552            break;
2553        }
2554    }
2555    let schema = Arc::new(ArrowSchema::new(vec![Field::new(
2556        "value",
2557        DataType::Int64,
2558        false,
2559    )]));
2560    let batch = arrow::record_batch::RecordBatch::try_new(
2561        schema.clone(),
2562        vec![Arc::new(Int64Array::from(values)) as ArrayRef],
2563    )?;
2564    Ok(Arc::new(MemTable::try_new(schema, vec![vec![batch]])?))
2565}
2566
2567#[derive(Clone, Copy)]
2568pub(crate) enum JsonTableMode {
2569    Each,
2570    Tree,
2571}
2572
2573fn json_table_function(
2574    args: TableFunctionArgs,
2575    name: &str,
2576    mode: JsonTableMode,
2577) -> DFResult<Arc<dyn TableProvider>> {
2578    let exprs = args.exprs();
2579    if exprs.is_empty() || exprs.len() > 2 {
2580        return Err(DataFusionError::Plan(format!(
2581            "{name} requires one JSON argument and an optional root path"
2582        )));
2583    }
2584    let json = literal_string(&exprs[0]).ok_or_else(|| {
2585        DataFusionError::Plan(format!("{name} arguments must be string literals"))
2586    })?;
2587    let root_path = if exprs.len() == 2 {
2588        literal_string(&exprs[1]).ok_or_else(|| {
2589            DataFusionError::Plan(format!("{name} root path must be a string literal"))
2590        })?
2591    } else {
2592        "$".to_string()
2593    };
2594    json_table_provider_from_text(name, &json, Some(&root_path), mode)
2595}
2596
2597#[derive(Default)]
2598struct JsonTableRows {
2599    key: Vec<Option<String>>,
2600    value: Vec<Option<String>>,
2601    ty: Vec<Option<String>>,
2602    atom: Vec<Option<String>>,
2603    id: Vec<i64>,
2604    parent: Vec<Option<i64>>,
2605    fullkey: Vec<Option<String>>,
2606    path: Vec<Option<String>>,
2607    next_id: i64,
2608}
2609
2610impl JsonTableRows {
2611    fn push(
2612        &mut self,
2613        row_key: Option<String>,
2614        row_value: &JsonValue,
2615        fullkey: String,
2616        path: String,
2617        parent: Option<i64>,
2618    ) -> i64 {
2619        let row_id = self.next_id;
2620        self.next_id += 1;
2621        self.key.push(row_key);
2622        self.value
2623            .push(json_scalar_result(row_value).or_else(|| Some(row_value.to_string())));
2624        self.ty.push(Some(json_type(row_value).to_string()));
2625        self.atom.push(match row_value {
2626            JsonValue::Array(_) | JsonValue::Object(_) => None,
2627            _ => json_scalar_result(row_value),
2628        });
2629        self.id.push(row_id);
2630        self.parent.push(parent);
2631        self.fullkey.push(Some(fullkey));
2632        self.path.push(Some(path));
2633        row_id
2634    }
2635}
2636
2637pub(crate) fn json_table_provider_from_text(
2638    name: &str,
2639    json: &str,
2640    root_path: Option<&str>,
2641    mode: JsonTableMode,
2642) -> DFResult<Arc<dyn TableProvider>> {
2643    let (schema, batches) = json_table_batches_from_text(name, json, root_path, mode)?;
2644    Ok(Arc::new(MemTable::try_new(schema, vec![batches])?))
2645}
2646
2647pub(crate) fn json_table_batches_from_text(
2648    name: &str,
2649    json: &str,
2650    root_path: Option<&str>,
2651    mode: JsonTableMode,
2652) -> DFResult<(Arc<ArrowSchema>, Vec<RecordBatch>)> {
2653    let root_path = root_path.unwrap_or("$");
2654    let root = serde_json::from_str::<JsonValue>(json)
2655        .map_err(|e| DataFusionError::Plan(format!("invalid JSON for {name}: {e}")))?;
2656    let path = parse_json_path(root_path)
2657        .ok_or_else(|| DataFusionError::Plan(format!("invalid {name} root path")))?;
2658    let value = json_path_get(&root, &path)
2659        .cloned()
2660        .unwrap_or(JsonValue::Null);
2661    json_table_batches(&value, root_path, mode)
2662}
2663
2664fn json_table_batches(
2665    value: &JsonValue,
2666    root_path: &str,
2667    mode: JsonTableMode,
2668) -> DFResult<(Arc<ArrowSchema>, Vec<RecordBatch>)> {
2669    let mut rows = JsonTableRows::default();
2670
2671    match mode {
2672        JsonTableMode::Each => json_each_rows(&mut rows, value, root_path),
2673        JsonTableMode::Tree => {
2674            json_tree_rows(
2675                &mut rows,
2676                None,
2677                value,
2678                root_path.to_string(),
2679                root_path.to_string(),
2680                None,
2681            );
2682        }
2683    }
2684
2685    let schema = Arc::new(ArrowSchema::new(vec![
2686        Field::new("key", DataType::Utf8, true),
2687        Field::new("value", DataType::Utf8, true),
2688        Field::new("type", DataType::Utf8, true),
2689        Field::new("atom", DataType::Utf8, true),
2690        Field::new("id", DataType::Int64, false),
2691        Field::new("parent", DataType::Int64, true),
2692        Field::new("fullkey", DataType::Utf8, true),
2693        Field::new("path", DataType::Utf8, true),
2694    ]));
2695    let batch = arrow::record_batch::RecordBatch::try_new(
2696        schema.clone(),
2697        vec![
2698            Arc::new(StringArray::from(rows.key)) as ArrayRef,
2699            Arc::new(StringArray::from(rows.value)),
2700            Arc::new(StringArray::from(rows.ty)),
2701            Arc::new(StringArray::from(rows.atom)),
2702            Arc::new(Int64Array::from(rows.id)),
2703            Arc::new(Int64Array::from(rows.parent)),
2704            Arc::new(StringArray::from(rows.fullkey)),
2705            Arc::new(StringArray::from(rows.path)),
2706        ],
2707    )?;
2708    Ok((schema, vec![batch]))
2709}
2710
2711fn json_each_rows(rows: &mut JsonTableRows, value: &JsonValue, root_path: &str) {
2712    match value {
2713        JsonValue::Array(values) => {
2714            for (idx, item) in values.iter().enumerate() {
2715                rows.push(
2716                    Some(idx.to_string()),
2717                    item,
2718                    format!("{root_path}[{idx}]"),
2719                    root_path.to_string(),
2720                    None,
2721                );
2722            }
2723        }
2724        JsonValue::Object(values) => {
2725            for (key, item) in values {
2726                rows.push(
2727                    Some(key.clone()),
2728                    item,
2729                    json_child_path(root_path, &JsonPathToken::Key(key.clone())),
2730                    root_path.to_string(),
2731                    None,
2732                );
2733            }
2734        }
2735        JsonValue::Null => {}
2736        scalar => {
2737            rows.push(
2738                None,
2739                scalar,
2740                root_path.to_string(),
2741                root_path.to_string(),
2742                None,
2743            );
2744        }
2745    }
2746}
2747
2748fn json_tree_rows(
2749    rows: &mut JsonTableRows,
2750    row_key: Option<String>,
2751    value: &JsonValue,
2752    fullkey: String,
2753    path: String,
2754    parent: Option<i64>,
2755) {
2756    let row_id = rows.push(row_key, value, fullkey.clone(), path, parent);
2757    match value {
2758        JsonValue::Array(values) => {
2759            for (idx, item) in values.iter().enumerate() {
2760                let child_fullkey = format!("{fullkey}[{idx}]");
2761                json_tree_rows(
2762                    rows,
2763                    Some(idx.to_string()),
2764                    item,
2765                    child_fullkey,
2766                    fullkey.clone(),
2767                    Some(row_id),
2768                );
2769            }
2770        }
2771        JsonValue::Object(values) => {
2772            for (key, item) in values {
2773                let token = JsonPathToken::Key(key.clone());
2774                let child_fullkey = json_child_path(&fullkey, &token);
2775                json_tree_rows(
2776                    rows,
2777                    Some(key.clone()),
2778                    item,
2779                    child_fullkey,
2780                    fullkey.clone(),
2781                    Some(row_id),
2782                );
2783            }
2784        }
2785        _ => {}
2786    }
2787}
2788
2789fn json_child_path(parent: &str, token: &JsonPathToken) -> String {
2790    match token {
2791        JsonPathToken::Index(idx) => format!("{parent}[{idx}]"),
2792        JsonPathToken::Key(key) if is_simple_json_key(key) => format!("{parent}.{key}"),
2793        JsonPathToken::Key(key) => format!("{parent}[{}]", JsonValue::String(key.clone())),
2794    }
2795}
2796
2797fn is_simple_json_key(key: &str) -> bool {
2798    let mut chars = key.chars();
2799    let Some(first) = chars.next() else {
2800        return false;
2801    };
2802    (first == '_' || first.is_ascii_alphabetic())
2803        && chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
2804}