partiql_logical_planner/
builtins.rs

1use itertools::Itertools;
2use once_cell::sync::Lazy;
3use partiql_catalog::call_defs::{CallDef, CallSpec, CallSpecArg};
4use partiql_logical as logical;
5use partiql_logical::{SetQuantifier, ValueExpr};
6use rustc_hash::FxHashMap;
7use std::fmt::Debug;
8use unicase::UniCase;
9
10fn function_call_def_char_len() -> CallDef {
11    CallDef {
12        names: vec!["char_length", "character_length"],
13        overloads: vec![CallSpec {
14            input: vec![CallSpecArg::Positional],
15            output: Box::new(|args| {
16                logical::ValueExpr::Call(logical::CallExpr {
17                    name: logical::CallName::CharLength,
18                    arguments: args,
19                })
20            }),
21        }],
22    }
23}
24
25fn function_call_def_octet_len() -> CallDef {
26    CallDef {
27        names: vec!["octet_length"],
28        overloads: vec![CallSpec {
29            input: vec![CallSpecArg::Positional],
30            output: Box::new(|args| {
31                logical::ValueExpr::Call(logical::CallExpr {
32                    name: logical::CallName::OctetLength,
33                    arguments: args,
34                })
35            }),
36        }],
37    }
38}
39
40fn function_call_def_bit_len() -> CallDef {
41    CallDef {
42        names: vec!["bit_length"],
43        overloads: vec![CallSpec {
44            input: vec![CallSpecArg::Positional],
45            output: Box::new(|args| {
46                logical::ValueExpr::Call(logical::CallExpr {
47                    name: logical::CallName::BitLength,
48                    arguments: args,
49                })
50            }),
51        }],
52    }
53}
54
55fn function_call_def_lower() -> CallDef {
56    CallDef {
57        names: vec!["lower"],
58        overloads: vec![CallSpec {
59            input: vec![CallSpecArg::Positional],
60            output: Box::new(|args| {
61                logical::ValueExpr::Call(logical::CallExpr {
62                    name: logical::CallName::Lower,
63                    arguments: args,
64                })
65            }),
66        }],
67    }
68}
69
70fn function_call_def_upper() -> CallDef {
71    CallDef {
72        names: vec!["upper"],
73        overloads: vec![CallSpec {
74            input: vec![CallSpecArg::Positional],
75            output: Box::new(|args| {
76                logical::ValueExpr::Call(logical::CallExpr {
77                    name: logical::CallName::Upper,
78                    arguments: args,
79                })
80            }),
81        }],
82    }
83}
84
85fn function_call_def_substring() -> CallDef {
86    CallDef {
87        names: vec!["substring"],
88        overloads: vec![
89            CallSpec {
90                input: vec![
91                    CallSpecArg::Positional,
92                    CallSpecArg::Positional,
93                    CallSpecArg::Positional,
94                ],
95                output: Box::new(|args| {
96                    logical::ValueExpr::Call(logical::CallExpr {
97                        name: logical::CallName::Substring,
98                        arguments: args,
99                    })
100                }),
101            },
102            CallSpec {
103                input: vec![CallSpecArg::Positional, CallSpecArg::Positional],
104                output: Box::new(|args| {
105                    logical::ValueExpr::Call(logical::CallExpr {
106                        name: logical::CallName::Substring,
107                        arguments: args,
108                    })
109                }),
110            },
111            CallSpec {
112                input: vec![
113                    CallSpecArg::Positional,
114                    CallSpecArg::Named("from".into()),
115                    CallSpecArg::Named("for".into()),
116                ],
117                output: Box::new(|args| {
118                    logical::ValueExpr::Call(logical::CallExpr {
119                        name: logical::CallName::Substring,
120                        arguments: args,
121                    })
122                }),
123            },
124            CallSpec {
125                input: vec![CallSpecArg::Positional, CallSpecArg::Named("from".into())],
126                output: Box::new(|args| {
127                    logical::ValueExpr::Call(logical::CallExpr {
128                        name: logical::CallName::Substring,
129                        arguments: args,
130                    })
131                }),
132            },
133            CallSpec {
134                input: vec![CallSpecArg::Positional, CallSpecArg::Named("for".into())],
135                output: Box::new(|mut args| {
136                    args.insert(1, ValueExpr::Lit(Box::new(logical::Lit::Int8(0))));
137                    logical::ValueExpr::Call(logical::CallExpr {
138                        name: logical::CallName::Substring,
139                        arguments: args,
140                    })
141                }),
142            },
143        ],
144    }
145}
146
147fn function_call_def_overlay() -> CallDef {
148    CallDef {
149        names: vec!["overlay"],
150        overloads: vec![
151            CallSpec {
152                input: vec![
153                    CallSpecArg::Positional,
154                    CallSpecArg::Named("placing".into()),
155                    CallSpecArg::Named("from".into()),
156                    CallSpecArg::Named("for".into()),
157                ],
158                output: Box::new(|args| {
159                    logical::ValueExpr::Call(logical::CallExpr {
160                        name: logical::CallName::Overlay,
161                        arguments: args,
162                    })
163                }),
164            },
165            CallSpec {
166                input: vec![
167                    CallSpecArg::Positional,
168                    CallSpecArg::Named("placing".into()),
169                    CallSpecArg::Named("from".into()),
170                ],
171                output: Box::new(|args| {
172                    logical::ValueExpr::Call(logical::CallExpr {
173                        name: logical::CallName::Overlay,
174                        arguments: args,
175                    })
176                }),
177            },
178        ],
179    }
180}
181
182fn function_call_def_position() -> CallDef {
183    CallDef {
184        names: vec!["position"],
185        overloads: vec![CallSpec {
186            input: vec![CallSpecArg::Positional, CallSpecArg::Named("in".into())],
187            output: Box::new(|args| {
188                logical::ValueExpr::Call(logical::CallExpr {
189                    name: logical::CallName::Position,
190                    arguments: args,
191                })
192            }),
193        }],
194    }
195}
196
197fn function_call_def_trim() -> CallDef {
198    CallDef {
199        names: vec!["trim"],
200        overloads: vec![
201            CallSpec {
202                input: vec![
203                    CallSpecArg::Named("leading".into()),
204                    CallSpecArg::Named("from".into()),
205                ],
206                output: Box::new(|args| {
207                    logical::ValueExpr::Call(logical::CallExpr {
208                        name: logical::CallName::LTrim,
209                        arguments: args,
210                    })
211                }),
212            },
213            CallSpec {
214                input: vec![
215                    CallSpecArg::Named("trailing".into()),
216                    CallSpecArg::Named("from".into()),
217                ],
218                output: Box::new(|args| {
219                    logical::ValueExpr::Call(logical::CallExpr {
220                        name: logical::CallName::RTrim,
221                        arguments: args,
222                    })
223                }),
224            },
225            CallSpec {
226                input: vec![
227                    CallSpecArg::Named("both".into()),
228                    CallSpecArg::Named("from".into()),
229                ],
230                output: Box::new(|args| {
231                    logical::ValueExpr::Call(logical::CallExpr {
232                        name: logical::CallName::BTrim,
233                        arguments: args,
234                    })
235                }),
236            },
237            CallSpec {
238                input: vec![CallSpecArg::Named("from".into())],
239                output: Box::new(|mut args| {
240                    args.insert(
241                        0,
242                        ValueExpr::Lit(Box::new(logical::Lit::String(" ".to_string()))),
243                    );
244
245                    logical::ValueExpr::Call(logical::CallExpr {
246                        name: logical::CallName::BTrim,
247                        arguments: args,
248                    })
249                }),
250            },
251            CallSpec {
252                input: vec![CallSpecArg::Positional],
253                output: Box::new(|mut args| {
254                    args.insert(
255                        0,
256                        ValueExpr::Lit(Box::new(logical::Lit::String(" ".to_string()))),
257                    );
258                    logical::ValueExpr::Call(logical::CallExpr {
259                        name: logical::CallName::BTrim,
260                        arguments: args,
261                    })
262                }),
263            },
264            CallSpec {
265                input: vec![CallSpecArg::Positional, CallSpecArg::Named("from".into())],
266                output: Box::new(|args| {
267                    logical::ValueExpr::Call(logical::CallExpr {
268                        name: logical::CallName::BTrim,
269                        arguments: args,
270                    })
271                }),
272            },
273        ],
274    }
275}
276
277fn function_call_def_coalesce() -> CallDef {
278    CallDef {
279        names: vec!["coalesce"],
280        overloads: (0..15)
281            .map(|n| CallSpec {
282                input: std::iter::repeat_n(CallSpecArg::Positional, n).collect_vec(),
283                output: Box::new(|args| {
284                    logical::ValueExpr::CoalesceExpr(logical::CoalesceExpr { elements: args })
285                }),
286            })
287            .collect_vec(),
288    }
289}
290
291fn function_call_def_nullif() -> CallDef {
292    CallDef {
293        names: vec!["nullif"],
294        overloads: vec![CallSpec {
295            input: vec![CallSpecArg::Positional, CallSpecArg::Positional],
296            output: Box::new(|mut args| {
297                assert_eq!(args.len(), 2);
298                let rhs = Box::new(args.pop().unwrap());
299                let lhs = Box::new(args.pop().unwrap());
300                logical::ValueExpr::NullIfExpr(logical::NullIfExpr { lhs, rhs })
301            }),
302        }],
303    }
304}
305
306fn function_call_def_exists() -> CallDef {
307    CallDef {
308        names: vec!["exists"],
309        overloads: vec![CallSpec {
310            input: vec![CallSpecArg::Positional],
311            output: Box::new(|args| {
312                logical::ValueExpr::Call(logical::CallExpr {
313                    name: logical::CallName::Exists,
314                    arguments: args,
315                })
316            }),
317        }],
318    }
319}
320
321fn function_call_def_abs() -> CallDef {
322    CallDef {
323        names: vec!["abs"],
324        overloads: vec![CallSpec {
325            input: vec![CallSpecArg::Positional],
326            output: Box::new(|args| {
327                logical::ValueExpr::Call(logical::CallExpr {
328                    name: logical::CallName::Abs,
329                    arguments: args,
330                })
331            }),
332        }],
333    }
334}
335
336fn function_call_def_mod() -> CallDef {
337    CallDef {
338        names: vec!["mod"],
339        overloads: vec![CallSpec {
340            input: vec![CallSpecArg::Positional, CallSpecArg::Positional],
341            output: Box::new(|args| {
342                logical::ValueExpr::Call(logical::CallExpr {
343                    name: logical::CallName::Mod,
344                    arguments: args,
345                })
346            }),
347        }],
348    }
349}
350
351fn function_call_def_cardinality() -> CallDef {
352    CallDef {
353        names: vec!["cardinality"],
354        overloads: vec![CallSpec {
355            input: vec![CallSpecArg::Positional],
356            output: Box::new(|args| {
357                logical::ValueExpr::Call(logical::CallExpr {
358                    name: logical::CallName::Cardinality,
359                    arguments: args,
360                })
361            }),
362        }],
363    }
364}
365
366fn function_call_def_extract() -> CallDef {
367    CallDef {
368        names: vec!["extract"],
369        overloads: vec![
370            CallSpec {
371                input: vec![
372                    CallSpecArg::Named("year".into()),
373                    CallSpecArg::Named("from".into()),
374                ],
375                output: Box::new(|mut args| {
376                    args.remove(0); // remove first default synthesized argument from parser preprocessor
377                    logical::ValueExpr::Call(logical::CallExpr {
378                        name: logical::CallName::ExtractYear,
379                        arguments: args,
380                    })
381                }),
382            },
383            CallSpec {
384                input: vec![
385                    CallSpecArg::Named("month".into()),
386                    CallSpecArg::Named("from".into()),
387                ],
388                output: Box::new(|mut args| {
389                    args.remove(0); // remove first default synthesized argument from parser preprocessor
390                    logical::ValueExpr::Call(logical::CallExpr {
391                        name: logical::CallName::ExtractMonth,
392                        arguments: args,
393                    })
394                }),
395            },
396            CallSpec {
397                input: vec![
398                    CallSpecArg::Named("day".into()),
399                    CallSpecArg::Named("from".into()),
400                ],
401                output: Box::new(|mut args| {
402                    args.remove(0); // remove first default synthesized argument from parser preprocessor
403                    logical::ValueExpr::Call(logical::CallExpr {
404                        name: logical::CallName::ExtractDay,
405                        arguments: args,
406                    })
407                }),
408            },
409            CallSpec {
410                input: vec![
411                    CallSpecArg::Named("hour".into()),
412                    CallSpecArg::Named("from".into()),
413                ],
414                output: Box::new(|mut args| {
415                    args.remove(0); // remove first default synthesized argument from parser preprocessor
416                    logical::ValueExpr::Call(logical::CallExpr {
417                        name: logical::CallName::ExtractHour,
418                        arguments: args,
419                    })
420                }),
421            },
422            CallSpec {
423                input: vec![
424                    CallSpecArg::Named("minute".into()),
425                    CallSpecArg::Named("from".into()),
426                ],
427                output: Box::new(|mut args| {
428                    args.remove(0); // remove first default synthesized argument from parser preprocessor
429                    logical::ValueExpr::Call(logical::CallExpr {
430                        name: logical::CallName::ExtractMinute,
431                        arguments: args,
432                    })
433                }),
434            },
435            CallSpec {
436                input: vec![
437                    CallSpecArg::Named("second".into()),
438                    CallSpecArg::Named("from".into()),
439                ],
440                output: Box::new(|mut args| {
441                    args.remove(0); // remove first default synthesized argument from parser preprocessor
442                    logical::ValueExpr::Call(logical::CallExpr {
443                        name: logical::CallName::ExtractSecond,
444                        arguments: args,
445                    })
446                }),
447            },
448            CallSpec {
449                input: vec![
450                    CallSpecArg::Named("timezone_hour".into()),
451                    CallSpecArg::Named("from".into()),
452                ],
453                output: Box::new(|mut args| {
454                    args.remove(0); // remove first default synthesized argument from parser preprocessor
455                    logical::ValueExpr::Call(logical::CallExpr {
456                        name: logical::CallName::ExtractTimezoneHour,
457                        arguments: args,
458                    })
459                }),
460            },
461            CallSpec {
462                input: vec![
463                    CallSpecArg::Named("timezone_minute".into()),
464                    CallSpecArg::Named("from".into()),
465                ],
466                output: Box::new(|mut args| {
467                    args.remove(0); // remove first default synthesized argument from parser preprocessor
468                    logical::ValueExpr::Call(logical::CallExpr {
469                        name: logical::CallName::ExtractTimezoneMinute,
470                        arguments: args,
471                    })
472                }),
473            },
474        ],
475    }
476}
477
478// ------------------------ COLL_AGG Functions ------------------------
479fn function_call_def_coll_avg() -> CallDef {
480    CallDef {
481        names: vec!["coll_avg"],
482        overloads: vec![
483            CallSpec {
484                input: vec![CallSpecArg::Positional],
485                output: Box::new(|args| {
486                    logical::ValueExpr::Call(logical::CallExpr {
487                        name: logical::CallName::CollAvg(SetQuantifier::All),
488                        arguments: args,
489                    })
490                }),
491            },
492            CallSpec {
493                input: vec![CallSpecArg::Named("all".into())],
494                output: Box::new(|args| {
495                    logical::ValueExpr::Call(logical::CallExpr {
496                        name: logical::CallName::CollAvg(SetQuantifier::All),
497                        arguments: args,
498                    })
499                }),
500            },
501            CallSpec {
502                input: vec![CallSpecArg::Named("distinct".into())],
503                output: Box::new(|args| {
504                    logical::ValueExpr::Call(logical::CallExpr {
505                        name: logical::CallName::CollAvg(SetQuantifier::Distinct),
506                        arguments: args,
507                    })
508                }),
509            },
510        ],
511    }
512}
513
514fn function_call_def_coll_count() -> CallDef {
515    CallDef {
516        names: vec!["coll_count"],
517        overloads: vec![
518            CallSpec {
519                input: vec![CallSpecArg::Positional],
520                output: Box::new(|args| {
521                    logical::ValueExpr::Call(logical::CallExpr {
522                        name: logical::CallName::CollCount(SetQuantifier::All),
523                        arguments: args,
524                    })
525                }),
526            },
527            CallSpec {
528                input: vec![CallSpecArg::Named("all".into())],
529                output: Box::new(|args| {
530                    logical::ValueExpr::Call(logical::CallExpr {
531                        name: logical::CallName::CollCount(SetQuantifier::All),
532                        arguments: args,
533                    })
534                }),
535            },
536            CallSpec {
537                input: vec![CallSpecArg::Named("distinct".into())],
538                output: Box::new(|args| {
539                    logical::ValueExpr::Call(logical::CallExpr {
540                        name: logical::CallName::CollCount(SetQuantifier::Distinct),
541                        arguments: args,
542                    })
543                }),
544            },
545        ],
546    }
547}
548
549fn function_call_def_coll_max() -> CallDef {
550    CallDef {
551        names: vec!["coll_max"],
552        overloads: vec![
553            CallSpec {
554                input: vec![CallSpecArg::Positional],
555                output: Box::new(|args| {
556                    logical::ValueExpr::Call(logical::CallExpr {
557                        name: logical::CallName::CollMax(SetQuantifier::All),
558                        arguments: args,
559                    })
560                }),
561            },
562            CallSpec {
563                input: vec![CallSpecArg::Named("all".into())],
564                output: Box::new(|args| {
565                    logical::ValueExpr::Call(logical::CallExpr {
566                        name: logical::CallName::CollMax(SetQuantifier::All),
567                        arguments: args,
568                    })
569                }),
570            },
571            CallSpec {
572                input: vec![CallSpecArg::Named("distinct".into())],
573                output: Box::new(|args| {
574                    logical::ValueExpr::Call(logical::CallExpr {
575                        name: logical::CallName::CollMax(SetQuantifier::Distinct),
576                        arguments: args,
577                    })
578                }),
579            },
580        ],
581    }
582}
583
584fn function_call_def_coll_min() -> CallDef {
585    CallDef {
586        names: vec!["coll_min"],
587        overloads: vec![
588            CallSpec {
589                input: vec![CallSpecArg::Positional],
590                output: Box::new(|args| {
591                    logical::ValueExpr::Call(logical::CallExpr {
592                        name: logical::CallName::CollMin(SetQuantifier::All),
593                        arguments: args,
594                    })
595                }),
596            },
597            CallSpec {
598                input: vec![CallSpecArg::Named("all".into())],
599                output: Box::new(|args| {
600                    logical::ValueExpr::Call(logical::CallExpr {
601                        name: logical::CallName::CollMin(SetQuantifier::All),
602                        arguments: args,
603                    })
604                }),
605            },
606            CallSpec {
607                input: vec![CallSpecArg::Named("distinct".into())],
608                output: Box::new(|args| {
609                    logical::ValueExpr::Call(logical::CallExpr {
610                        name: logical::CallName::CollMin(SetQuantifier::Distinct),
611                        arguments: args,
612                    })
613                }),
614            },
615        ],
616    }
617}
618
619fn function_call_def_coll_sum() -> CallDef {
620    CallDef {
621        names: vec!["coll_sum"],
622        overloads: vec![
623            CallSpec {
624                input: vec![CallSpecArg::Positional],
625                output: Box::new(|args| {
626                    logical::ValueExpr::Call(logical::CallExpr {
627                        name: logical::CallName::CollSum(SetQuantifier::All),
628                        arguments: args,
629                    })
630                }),
631            },
632            CallSpec {
633                input: vec![CallSpecArg::Named("all".into())],
634                output: Box::new(|args| {
635                    logical::ValueExpr::Call(logical::CallExpr {
636                        name: logical::CallName::CollSum(SetQuantifier::All),
637                        arguments: args,
638                    })
639                }),
640            },
641            CallSpec {
642                input: vec![CallSpecArg::Named("distinct".into())],
643                output: Box::new(|args| {
644                    logical::ValueExpr::Call(logical::CallExpr {
645                        name: logical::CallName::CollSum(SetQuantifier::Distinct),
646                        arguments: args,
647                    })
648                }),
649            },
650        ],
651    }
652}
653
654fn function_call_def_coll_any() -> CallDef {
655    CallDef {
656        names: vec!["coll_any", "coll_some"],
657        overloads: vec![
658            CallSpec {
659                input: vec![CallSpecArg::Positional],
660                output: Box::new(|args| {
661                    logical::ValueExpr::Call(logical::CallExpr {
662                        name: logical::CallName::CollAny(SetQuantifier::All),
663                        arguments: args,
664                    })
665                }),
666            },
667            CallSpec {
668                input: vec![CallSpecArg::Named("all".into())],
669                output: Box::new(|args| {
670                    logical::ValueExpr::Call(logical::CallExpr {
671                        name: logical::CallName::CollAny(SetQuantifier::All),
672                        arguments: args,
673                    })
674                }),
675            },
676            CallSpec {
677                input: vec![CallSpecArg::Named("distinct".into())],
678                output: Box::new(|args| {
679                    logical::ValueExpr::Call(logical::CallExpr {
680                        name: logical::CallName::CollAny(SetQuantifier::Distinct),
681                        arguments: args,
682                    })
683                }),
684            },
685        ],
686    }
687}
688
689fn function_call_def_coll_every() -> CallDef {
690    CallDef {
691        names: vec!["coll_every"],
692        overloads: vec![
693            CallSpec {
694                input: vec![CallSpecArg::Positional],
695                output: Box::new(|args| {
696                    logical::ValueExpr::Call(logical::CallExpr {
697                        name: logical::CallName::CollEvery(SetQuantifier::All),
698                        arguments: args,
699                    })
700                }),
701            },
702            CallSpec {
703                input: vec![CallSpecArg::Named("all".into())],
704                output: Box::new(|args| {
705                    logical::ValueExpr::Call(logical::CallExpr {
706                        name: logical::CallName::CollEvery(SetQuantifier::All),
707                        arguments: args,
708                    })
709                }),
710            },
711            CallSpec {
712                input: vec![CallSpecArg::Named("distinct".into())],
713                output: Box::new(|args| {
714                    logical::ValueExpr::Call(logical::CallExpr {
715                        name: logical::CallName::CollEvery(SetQuantifier::Distinct),
716                        arguments: args,
717                    })
718                }),
719            },
720        ],
721    }
722}
723
724pub(crate) static FN_SYM_TAB: Lazy<FnSymTab> = Lazy::new(function_call_def);
725
726/// Function symbol table
727#[derive(Debug)]
728pub struct FnSymTab {
729    calls: FxHashMap<UniCase<String>, CallDef>,
730    synonyms: FxHashMap<UniCase<String>, UniCase<String>>,
731}
732
733impl FnSymTab {
734    pub fn lookup(&self, fn_name: &str) -> Option<&CallDef> {
735        self.synonyms
736            .get(&fn_name.into())
737            .and_then(|name| self.calls.get(name))
738    }
739}
740
741pub fn function_call_def() -> FnSymTab {
742    let mut calls: FxHashMap<UniCase<String>, CallDef> = FxHashMap::default();
743    let mut synonyms: FxHashMap<UniCase<String>, UniCase<String>> = FxHashMap::default();
744
745    for def in [
746        function_call_def_char_len(),
747        function_call_def_octet_len(),
748        function_call_def_bit_len(),
749        function_call_def_lower(),
750        function_call_def_upper(),
751        function_call_def_substring(),
752        function_call_def_position(),
753        function_call_def_overlay(),
754        function_call_def_trim(),
755        function_call_def_coalesce(),
756        function_call_def_nullif(),
757        function_call_def_exists(),
758        function_call_def_abs(),
759        function_call_def_mod(),
760        function_call_def_cardinality(),
761        function_call_def_extract(),
762        function_call_def_coll_avg(),
763        function_call_def_coll_count(),
764        function_call_def_coll_max(),
765        function_call_def_coll_min(),
766        function_call_def_coll_sum(),
767        function_call_def_coll_any(),
768        function_call_def_coll_every(),
769    ] {
770        assert!(!def.names.is_empty());
771        let primary = def.names[0];
772        synonyms.insert(primary.into(), primary.into());
773        for &name in &def.names[1..] {
774            synonyms.insert(name.into(), primary.into());
775        }
776
777        calls.insert(primary.into(), def);
778    }
779
780    FnSymTab { calls, synonyms }
781}