vortex_array/scalar_fn/
foreign.rs1use 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#[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#[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}