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