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 "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 "sqlite_compileoption_get",
160 "sqlite_compileoption_used",
161 "sqlite_offset",
162 "sqlite_source_id",
163 "sqlite_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 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 "sqlite_compileoption_get",
892 CompileOptionGet,
893 Volatility::Immutable,
894 ),
895 ExtendedSqlFunc::new(
896 "sqlite_compileoption_used",
897 CompileOptionUsed,
898 Volatility::Immutable,
899 ),
900 ExtendedSqlFunc::new("sqlite_offset", Offset, Volatility::Immutable),
901 ExtendedSqlFunc::nullary("sqlite_source_id", SourceId, Volatility::Immutable),
902 ExtendedSqlFunc::nullary("sqlite_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}