partiql_extension_value_functions/
lib.rs

1#![deny(rust_2018_idioms)]
2#![deny(clippy::all)]
3
4use partiql_catalog::call_defs::ScalarFnCallDef;
5use partiql_catalog::catalog::Catalog;
6use partiql_catalog::context::SessionContext;
7use partiql_catalog::extension::ExtensionResultError;
8use partiql_catalog::scalar_fn::{
9    vararg_scalar_fn_overloads, ScalarFnExpr, ScalarFnExprResult, ScalarFunction,
10    SimpleScalarFunctionInfo,
11};
12use partiql_value::{Tuple, Value};
13use std::borrow::Cow;
14
15#[derive(Debug, Default)]
16pub struct PartiqlValueFnExtension {}
17
18impl partiql_catalog::extension::Extension for PartiqlValueFnExtension {
19    fn name(&self) -> String {
20        "value-functions".into()
21    }
22
23    fn load(&self, catalog: &mut dyn Catalog) -> Result<(), ExtensionResultError> {
24        for scfn in [function_catalog_tupleunion, function_catalog_tupleconcat] {
25            match catalog.add_scalar_function(scfn()) {
26                Ok(_) => continue,
27                Err(e) => return Err(ExtensionResultError::LoadError(e.into())),
28            }
29        }
30        Ok(())
31    }
32}
33
34fn function_catalog_tupleunion() -> ScalarFunction {
35    let scalar_fn = Box::new(TupleUnionFnExpr::default());
36    let call_def = ScalarFnCallDef {
37        names: vec!["tupleunion"],
38        overloads: vararg_scalar_fn_overloads(scalar_fn),
39    };
40
41    let info = SimpleScalarFunctionInfo::new(call_def);
42    ScalarFunction::new(Box::new(info))
43}
44
45/// Represents a built-in tupleunion function,
46/// e.g. `tupleunion({ 'bob': 1 }, { 'sally': 2 }, { 'sally': 2 })` -> `{'bob: 1, 'sally':1, 'sally':2}`.
47#[derive(Debug, Clone, Default)]
48struct TupleUnionFnExpr {}
49impl ScalarFnExpr for TupleUnionFnExpr {
50    fn evaluate<'c>(
51        &self,
52        args: &[Cow<'_, Value>],
53        _ctx: &'c dyn SessionContext,
54    ) -> ScalarFnExprResult<'c> {
55        let mut t = Tuple::default();
56        for arg in args {
57            t.extend(
58                arg.as_tuple_ref()
59                    .pairs()
60                    .map(|(k, v)| (k.as_str(), v.clone())),
61            )
62        }
63        Ok(Cow::Owned(Value::from(t)))
64    }
65}
66
67fn function_catalog_tupleconcat() -> ScalarFunction {
68    let scalar_fn = Box::new(TupleConcatFnExpr::default());
69    let call_def = ScalarFnCallDef {
70        names: vec!["tupleconcat"],
71        overloads: vararg_scalar_fn_overloads(scalar_fn),
72    };
73
74    let info = SimpleScalarFunctionInfo::new(call_def);
75    ScalarFunction::new(Box::new(info))
76}
77
78/// Represents a built-in tupleconcat function,
79/// e.g. `tupleconcat({ 'bob': 1 }, { 'sally': 2 }, { 'sally': 2 })` -> `{'bob: 1, 'sally':2}`.
80#[derive(Debug, Clone, Default)]
81struct TupleConcatFnExpr {}
82impl ScalarFnExpr for TupleConcatFnExpr {
83    fn evaluate<'c>(
84        &self,
85        args: &[Cow<'_, Value>],
86        _ctx: &'c dyn SessionContext,
87    ) -> ScalarFnExprResult<'c> {
88        let result = args
89            .iter()
90            .map(|val| val.as_tuple_ref())
91            .reduce(|l, r| Cow::Owned(l.tuple_concat(&r)))
92            .map(|v| v.into_owned())
93            .unwrap_or_default();
94        Ok(Cow::Owned(Value::from(result)))
95    }
96}