Skip to main content

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::ScalarFnId;
21use crate::scalar_fn::ScalarFnRef;
22use crate::scalar_fn::ScalarFnVTable;
23use crate::scalar_fn::TypedScalarFnInstance;
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        TypedScalarFnInstance::new(Self::new(id), ForeignScalarFnOptions::new(metadata, arity))
62            .erased()
63    }
64}
65
66impl ScalarFnVTable for ForeignScalarFnVTable {
67    type Options = ForeignScalarFnOptions;
68
69    fn id(&self) -> ScalarFnId {
70        self.id
71    }
72
73    fn serialize(&self, options: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
74        Ok(Some(options.metadata.clone()))
75    }
76
77    fn deserialize(
78        &self,
79        metadata: &[u8],
80        _session: &VortexSession,
81    ) -> VortexResult<Self::Options> {
82        Ok(ForeignScalarFnOptions::new(metadata.to_vec(), 0))
83    }
84
85    fn arity(&self, options: &Self::Options) -> Arity {
86        Arity::Exact(options.arity)
87    }
88
89    fn child_name(&self, _options: &Self::Options, child_idx: usize) -> ChildName {
90        ChildName::new_arc(format!("arg{child_idx}").into())
91    }
92
93    fn fmt_sql(
94        &self,
95        _options: &Self::Options,
96        expr: &Expression,
97        f: &mut Formatter<'_>,
98    ) -> fmt::Result {
99        write!(f, "{}(", self.id)?;
100        for i in 0..expr.children().len() {
101            if i > 0 {
102                write!(f, ", ")?;
103            }
104            expr.child(i).fmt_sql(f)?;
105        }
106        write!(f, ")")
107    }
108
109    fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult<DType> {
110        Ok(DType::Variant(Nullability::Nullable))
111    }
112
113    fn execute(
114        &self,
115        _options: &Self::Options,
116        _args: &dyn ExecutionArgs,
117        _ctx: &mut ExecutionCtx,
118    ) -> VortexResult<ArrayRef> {
119        vortex_bail!("Cannot execute unknown scalar function '{}'", self.id);
120    }
121}