risingwave_expr/sig/
udf.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! UDF implementation interface.
16//!
17//! To support a new language or runtime for UDF, implement the interface in this module.
18//!
19//! See expr/impl/src/udf for the implementations.
20
21use anyhow::{Context, Result, bail};
22use educe::Educe;
23use enum_as_inner::EnumAsInner;
24use futures::stream::BoxStream;
25use risingwave_common::array::arrow::arrow_array_udf::{ArrayRef, BooleanArray, RecordBatch};
26use risingwave_common::types::DataType;
27
28/// The global registry of UDF implementations.
29///
30/// To register a new UDF implementation:
31///
32/// ```ignore
33/// #[linkme::distributed_slice(UDF_IMPLS)]
34/// static MY_UDF_LANGUAGE: UdfImplDescriptor = UdfImplDescriptor {...};
35/// ```
36#[linkme::distributed_slice]
37pub static UDF_IMPLS: [UdfImplDescriptor];
38
39/// Find a UDF implementation by language.
40pub fn find_udf_impl(
41    language: &str,
42    runtime: Option<&str>,
43    link: Option<&str>,
44) -> Result<&'static UdfImplDescriptor> {
45    let mut impls = UDF_IMPLS
46        .iter()
47        .filter(|desc| (desc.match_fn)(language, runtime, link));
48    let impl_ = impls.next().context(
49        "language not found.\nHINT: UDF feature flag may not be enabled during compilation",
50    )?;
51    if impls.next().is_some() {
52        bail!("multiple UDF implementations found for language: {language}");
53    }
54    Ok(impl_)
55}
56
57/// UDF implementation descriptor.
58///
59/// Every UDF implementation describes how it is selected, created, and built, together with the
60/// capabilities supported by its execution path.
61pub struct UdfImplDescriptor {
62    /// Returns if a function matches the implementation.
63    ///
64    /// This function is used to determine which implementation to use for a UDF.
65    pub match_fn: fn(language: &str, runtime: Option<&str>, link: Option<&str>) -> bool,
66
67    /// Creates a function from options.
68    ///
69    /// This function will be called when `create function` statement is executed on the frontend.
70    pub create_fn: fn(opts: CreateOptions<'_>) -> Result<CreateFunctionOutput>,
71
72    /// Builds UDF runtime from verified options.
73    ///
74    /// This function will be called before the UDF is executed on the backend.
75    pub build_fn: fn(opts: BuildOptions<'_>) -> Result<Box<dyn UdfImpl>>,
76
77    /// Returns whether this implementation and function kind honor
78    /// `always_retry_on_network_error` during execution.
79    pub supports_always_retry_on_network_error: fn(kind: UdfKind) -> bool,
80}
81
82/// Options for creating a function.
83///
84/// These information are parsed from `CREATE FUNCTION` statement.
85/// Implementations should verify the options and return a `CreateFunctionOutput` in `create_fn`.
86pub struct CreateOptions<'a> {
87    pub kind: UdfKind,
88    /// The function name registered in RisingWave.
89    pub name: &'a str,
90    pub arg_names: &'a [String],
91    pub arg_types: &'a [DataType],
92    pub return_type: &'a DataType,
93    /// The function name on the remote side / in the source code, currently only used for external UDF.
94    pub as_: Option<&'a str>,
95    pub using_link: Option<&'a str>,
96    pub using_base64_decoded: Option<&'a [u8]>,
97}
98
99/// Output of creating a function.
100pub struct CreateFunctionOutput {
101    /// The name for identifying the function in the UDF runtime.
102    pub name_in_runtime: String,
103    pub body: Option<String>,
104    pub compressed_binary: Option<Vec<u8>>,
105}
106
107/// Options for building a UDF runtime.
108#[derive(Educe)]
109#[educe(Debug)]
110pub struct BuildOptions<'a> {
111    pub kind: UdfKind,
112    pub body: Option<&'a str>,
113    #[educe(Debug(ignore))]
114    pub compressed_binary: Option<&'a [u8]>,
115    pub link: Option<&'a str>,
116    pub name_in_runtime: &'a str,
117    pub arg_names: &'a [String],
118    pub arg_types: &'a [DataType],
119    pub return_type: &'a DataType,
120    pub always_retry_on_network_error: bool,
121    pub language: &'a str,
122    pub is_async: Option<bool>,
123    pub is_batched: Option<bool>,
124}
125
126#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, EnumAsInner)]
127pub enum UdfKind {
128    Scalar,
129    Table,
130    Aggregate,
131}
132
133/// UDF implementation.
134#[async_trait::async_trait]
135pub trait UdfImpl: std::fmt::Debug + Send + Sync {
136    /// Call the scalar function.
137    async fn call(&self, input: &RecordBatch) -> Result<RecordBatch>;
138
139    /// Call the table function.
140    async fn call_table_function<'a>(
141        &'a self,
142        input: &'a RecordBatch,
143    ) -> Result<BoxStream<'a, Result<RecordBatch>>>;
144
145    /// For aggregate function, create the initial state.
146    async fn call_agg_create_state(&self) -> Result<ArrayRef> {
147        bail!("aggregate function is not supported");
148    }
149
150    /// For aggregate function, accumulate or retract the state.
151    async fn call_agg_accumulate_or_retract(
152        &self,
153        _state: &ArrayRef,
154        _ops: &BooleanArray,
155        _input: &RecordBatch,
156    ) -> Result<ArrayRef> {
157        bail!("aggregate function is not supported");
158    }
159
160    /// For aggregate function, get aggregate result from the state.
161    async fn call_agg_finish(&self, _state: &ArrayRef) -> Result<ArrayRef> {
162        bail!("aggregate function is not supported");
163    }
164
165    /// Whether the UDF talks in legacy mode.
166    ///
167    /// If true, decimal and jsonb types are mapped to Arrow `LargeBinary` and `LargeUtf8` types.
168    /// Otherwise, they are mapped to Arrow extension types.
169    /// See <https://github.com/risingwavelabs/arrow-udf/tree/main#extension-types>.
170    fn is_legacy(&self) -> bool {
171        false
172    }
173
174    /// Return the memory size consumed by UDF runtime in bytes.
175    ///
176    /// If not available, return 0.
177    fn memory_usage(&self) -> usize {
178        0
179    }
180}
risingwave_expr/sig/udf.rs

risingwave_expr/sig/
udf.rs
