vortex_array/scalar_fn/
foreign.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt;
5use std::fmt::Display;
6use std::fmt::Formatter;
7
8use vortex_error::VortexResult;
9use vortex_error::vortex_bail;
10use vortex_session::VortexSession;
11
12use crate::ArrayRef;
13use crate::ExecutionCtx;
14use crate::dtype::DType;
15use crate::dtype::Nullability;
16use crate::expr::Expression;
17use crate::scalar_fn::Arity;
18use crate::scalar_fn::ChildName;
19use crate::scalar_fn::ExecutionArgs;
20use crate::scalar_fn::ScalarFn;
21use crate::scalar_fn::ScalarFnId;
22use crate::scalar_fn::ScalarFnRef;
23use crate::scalar_fn::ScalarFnVTable;
24
25/// Options payload for a foreign scalar function.
26#[derive(Clone, Debug, PartialEq, Eq, Hash)]
27pub struct ForeignScalarFnOptions {
28    metadata: Vec<u8>,
29    arity: usize,
30}
31
32impl ForeignScalarFnOptions {
33    pub fn new(metadata: Vec<u8>, arity: usize) -> Self {
34        Self { metadata, arity }
35    }
36}
37
38impl Display for ForeignScalarFnOptions {
39    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
40        write!(
41            f,
42            "foreign(arity={}, metadata={}B)",
43            self.arity,
44            self.metadata.len()
45        )
46    }
47}
48
49/// Scalar function placeholder used when deserializing an unknown scalar function ID.
50#[derive(Clone, Debug)]
51pub struct ForeignScalarFnVTable {
52    id: ScalarFnId,
53}
54
55impl ForeignScalarFnVTable {
56    pub fn new(id: ScalarFnId) -> Self {
57        Self { id }
58    }
59
60    pub fn make_scalar_fn(id: ScalarFnId, metadata: Vec<u8>, arity: usize) -> ScalarFnRef {
61        ScalarFn::new(Self::new(id), ForeignScalarFnOptions::new(metadata, arity)).erased()
62    }
63}
64
65impl ScalarFnVTable for ForeignScalarFnVTable {
66    type Options = ForeignScalarFnOptions;
67
68    fn id(&self) -> ScalarFnId {
69        self.id.clone()
70    }
71
72    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
73        Ok(Some(options.metadata.clone()))
74    }
75
76    fn deserialize(
77        &self,
78        metadata: &[u8],
79        _session: &VortexSession,
80    ) -> VortexResult<Self::Options> {
81        Ok(ForeignScalarFnOptions::new(metadata.to_vec(), 0))
82    }
83
84    fn arity(&self, options: &Self::Options) -> Arity {
85        Arity::Exact(options.arity)
86    }
87
88    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
89        ChildName::new_arc(format!("arg{child_idx}").into())
90    }
91
92    fn fmt_sql(
93        &self,
94        _options: &Self::Options,
95        expr: &Expression,
96        f: &mut Formatter<'_>,
97    ) -> fmt::Result {
98        write!(f, "{}(", self.id)?;
99        for i in 0..expr.children().len() {
100            if i > 0 {
101                write!(f, ", ")?;
102            }
103            expr.child(i).fmt_sql(f)?;
104        }
105        write!(f, ")")
106    }
107
108    fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult<DType> {
109        Ok(DType::Variant(Nullability::Nullable))
110    }
111
112    fn execute(
113        &self,
114        _options: &Self::Options,
115        _args: &dyn ExecutionArgs,
116        _ctx: &mut ExecutionCtx,
117    ) -> VortexResult<ArrayRef> {
118        vortex_bail!("Cannot execute unknown scalar function '{}'", self.id);
119    }
120}
vortex_array/scalar_fn/foreign.rs

vortex_array/scalar_fn/
foreign.rs
