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