risingwave_expr/aggregate/
user_defined.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
15use std::sync::Arc;
16
17use anyhow::Context;
18use risingwave_common::array::Op;
19use risingwave_common::array::arrow::arrow_array_udf::ArrayRef;
20use risingwave_common::array::arrow::arrow_schema_udf::{Field, Fields, Schema, SchemaRef};
21use risingwave_common::array::arrow::{UdfArrowConvert, UdfFromArrow, UdfToArrow};
22use risingwave_common::bitmap::Bitmap;
23use risingwave_pb::expr::PbUserDefinedFunctionMetadata;
24
25use super::*;
26use crate::sig::{BuildOptions, UdfImpl, UdfKind};
27
28#[derive(Debug)]
29pub struct UserDefinedAggregateFunction {
30    arg_schema: SchemaRef,
31    return_type: DataType,
32    return_field: Field,
33    state_field: Field,
34    runtime: Box<dyn UdfImpl>,
35}
36
37#[async_trait::async_trait]
38impl AggregateFunction for UserDefinedAggregateFunction {
39    fn return_type(&self) -> DataType {
40        self.return_type.clone()
41    }
42
43    /// Creates an initial state of the aggregate function.
44    fn create_state(&self) -> Result<AggregateState> {
45        // FIXME(eric): This is bad. Let's make `create_state` async if someday we allow async UDAF
46        futures::executor::block_on(async {
47            let state = self.runtime.call_agg_create_state().await?;
48            Ok(AggregateState::Any(Box::new(State(state))))
49        })
50    }
51
52    /// Update the state with multiple rows.
53    async fn update(&self, state: &mut AggregateState, input: &StreamChunk) -> Result<()> {
54        let state = &mut state.downcast_mut::<State>().0;
55        let ops = input
56            .visibility()
57            .iter_ones()
58            .map(|i| Some(matches!(input.ops()[i], Op::Delete | Op::UpdateDelete)))
59            .collect();
60        // this will drop invisible rows
61        let arrow_input = UdfArrowConvert::default()
62            .to_record_batch(self.arg_schema.clone(), input.data_chunk())?;
63        let new_state = self
64            .runtime
65            .call_agg_accumulate_or_retract(state, &ops, &arrow_input)
66            .await?;
67        *state = new_state;
68        Ok(())
69    }
70
71    /// Update the state with a range of rows.
72    async fn update_range(
73        &self,
74        state: &mut AggregateState,
75        input: &StreamChunk,
76        range: Range<usize>,
77    ) -> Result<()> {
78        // XXX(runji): this may be inefficient
79        let vis = input.visibility() & Bitmap::from_range(input.capacity(), range);
80        let input = input.clone_with_vis(vis);
81        self.update(state, &input).await
82    }
83
84    /// Get aggregate result from the state.
85    async fn get_result(&self, state: &AggregateState) -> Result<Datum> {
86        let state = &state.downcast_ref::<State>().0;
87        let arrow_output = self.runtime.call_agg_finish(state).await?;
88        let output = UdfArrowConvert::default().from_array(&self.return_field, &arrow_output)?;
89        // The UDF runtime is external input: a server may drift from the signature it was
90        // checked against at creation time. Surface a mistyped result instead of letting it
91        // corrupt downstream value encoding.
92        if output.data_type() != self.return_type {
93            return Err(anyhow::anyhow!(
94                "UDF returned a value of type {} while the declared return type is {}",
95                output.data_type(),
96                self.return_type
97            )
98            .into());
99        }
100        Ok(output.datum_at(0))
101    }
102
103    /// Encode the state into a datum that can be stored in state table.
104    fn encode_state(&self, state: &AggregateState) -> Result<Datum> {
105        let state = &state.downcast_ref::<State>().0;
106        let state = UdfArrowConvert::default().from_array(&self.state_field, state)?;
107        Ok(state.datum_at(0))
108    }
109
110    /// Decode the state from a datum in state table.
111    fn decode_state(&self, datum: Datum) -> Result<AggregateState> {
112        let array = {
113            let mut builder = DataType::Bytea.create_array_builder(1);
114            builder.append(datum);
115            builder.finish()
116        };
117        let state = UdfArrowConvert::default().to_array(self.state_field.data_type(), &array)?;
118        Ok(AggregateState::Any(Box::new(State(state))))
119    }
120}
121
122// In arrow-udf, aggregate state is represented as an `ArrayRef`.
123// To avoid unnecessary conversion between `ArrayRef` and `Datum`,
124// we store `ArrayRef` directly in our `AggregateState`.
125#[derive(Debug)]
126struct State(ArrayRef);
127
128impl EstimateSize for State {
129    fn estimated_heap_size(&self) -> usize {
130        self.0.get_array_memory_size()
131    }
132}
133
134impl AggStateDyn for State {}
135
136/// Create a new user-defined aggregate function.
137pub fn new_user_defined(
138    return_type: &DataType,
139    udf: &PbUserDefinedFunctionMetadata,
140) -> Result<BoxedAggregateFunction> {
141    let arg_types = udf.arg_types.iter().map(|t| t.into()).collect::<Vec<_>>();
142    let language = udf.language.as_str();
143    let runtime = udf.runtime.as_deref();
144    let link = udf.link.as_deref();
145
146    let name_in_runtime = udf
147        .name_in_runtime()
148        .expect("SQL UDF won't get here, other UDFs must have `name_in_runtime`");
149
150    let build_fn = crate::sig::find_udf_impl(language, runtime, link)?.build_fn;
151    let runtime = build_fn(BuildOptions {
152        kind: UdfKind::Aggregate,
153        body: udf.body.as_deref(),
154        compressed_binary: udf.compressed_binary.as_deref(),
155        link: udf.link.as_deref(),
156        name_in_runtime,
157        arg_names: &udf.arg_names,
158        arg_types: &arg_types,
159        return_type,
160        always_retry_on_network_error: false,
161        language,
162        is_async: udf.is_async,
163        is_batched: udf.is_batched,
164    })
165    .context("failed to build UDF runtime")?;
166
167    // legacy UDF runtimes do not support aggregate functions,
168    // so we can assume that the runtime is not legacy
169    let arrow_convert = UdfArrowConvert::default();
170    let arg_schema = Arc::new(Schema::new(
171        arg_types
172            .iter()
173            .map(|t| arrow_convert.to_arrow_field("", t))
174            .try_collect::<_, Fields, _>()?,
175    ));
176
177    Ok(Box::new(UserDefinedAggregateFunction {
178        return_field: arrow_convert.to_arrow_field("", return_type)?,
179        state_field: Field::new(
180            "state",
181            risingwave_common::array::arrow::arrow_schema_udf::DataType::Binary,
182            true,
183        ),
184        return_type: return_type.clone(),
185        arg_schema,
186        runtime,
187    }))
188}
189
190#[cfg(test)]
191mod tests {
192    use anyhow::Result;
193    use futures::stream::BoxStream;
194    use risingwave_common::array::arrow::arrow_array_udf::{
195        Float64Array, Int32Array, Int64Array, MapArray, RecordBatch, StringArray, StructArray,
196    };
197    use risingwave_common::array::arrow::arrow_buffer_udf::OffsetBuffer;
198    use risingwave_common::array::arrow::arrow_schema_udf::DataType as ArrowDataType;
199    use risingwave_common::types::{MapType, StructType};
200
201    use super::*;
202    use crate::sig::UdfImpl;
203
204    /// Returns the given array from `finish`, regardless of the declared return type.
205    #[derive(Debug)]
206    struct MistypedRuntime(ArrayRef);
207
208    #[async_trait::async_trait]
209    impl UdfImpl for MistypedRuntime {
210        async fn call(&self, _input: &RecordBatch) -> Result<RecordBatch> {
211            unimplemented!()
212        }
213
214        async fn call_table_function<'a>(
215            &'a self,
216            _input: &'a RecordBatch,
217        ) -> Result<BoxStream<'a, Result<RecordBatch>>> {
218            unimplemented!()
219        }
220
221        async fn call_agg_finish(&self, _state: &ArrayRef) -> Result<ArrayRef> {
222            Ok(self.0.clone())
223        }
224    }
225
226    /// Drives `get_result` with a runtime that returns `output` and yields the error message.
227    async fn get_result_err(return_type: DataType, output: ArrayRef) -> String {
228        let convert = UdfArrowConvert::default();
229        let agg = UserDefinedAggregateFunction {
230            return_field: convert.to_arrow_field("", &return_type).unwrap(),
231            state_field: Field::new("state", ArrowDataType::Binary, true),
232            return_type,
233            arg_schema: Arc::new(Schema::new(Vec::<Field>::new())),
234            runtime: Box::new(MistypedRuntime(output)),
235        };
236        let state = AggregateState::Any(Box::new(State(
237            Arc::new(Int64Array::from(vec![0i64])) as ArrayRef
238        )));
239        let err = agg.get_result(&state).await.unwrap_err();
240        format!("{:?}", anyhow::anyhow!(err))
241    }
242
243    #[tokio::test]
244    async fn misbehaving_runtime_mistyped_finish() {
245        // Struct child diverges from the declared `struct<a bigint>`.
246        let fields: Fields = vec![Field::new("a", ArrowDataType::Utf8, true)].into();
247        let struct_output: ArrayRef = Arc::new(StructArray::new(
248            fields,
249            vec![Arc::new(StringArray::from(vec![Some("oops")])) as ArrayRef],
250            None,
251        ));
252        let msg = get_result_err(
253            DataType::Struct(StructType::new(vec![("a", DataType::Int64)])),
254            struct_output,
255        )
256        .await;
257        assert!(
258            msg.contains("declared return type"),
259            "unexpected error: {msg}"
260        );
261
262        // Scalar output diverges from the declared `bigint`.
263        let msg = get_result_err(DataType::Int64, Arc::new(Int32Array::from(vec![Some(1)]))).await;
264        assert!(
265            msg.contains("declared return type"),
266            "unexpected error: {msg}"
267        );
268
269        // Map key type is unrepresentable in RW: must be a graceful error, not a panic
270        // from `MapArray::data_type()`.
271        let entries_fields: Fields = vec![
272            Field::new("key", ArrowDataType::Float64, false),
273            Field::new("value", ArrowDataType::Int32, true),
274        ]
275        .into();
276        let entries = StructArray::new(
277            entries_fields.clone(),
278            vec![
279                Arc::new(Float64Array::from(vec![Some(1.5)])) as ArrayRef,
280                Arc::new(Int32Array::from(vec![Some(42)])),
281            ],
282            None,
283        );
284        let map_output: ArrayRef = Arc::new(MapArray::new(
285            Arc::new(Field::new(
286                "entries",
287                ArrowDataType::Struct(entries_fields),
288                false,
289            )),
290            OffsetBuffer::new(vec![0, 1].into()),
291            entries,
292            None,
293            false,
294        ));
295        let msg = get_result_err(
296            DataType::Map(MapType::from_kv(DataType::Varchar, DataType::Int32)),
297            map_output,
298        )
299        .await;
300        assert!(
301            msg.contains("invalid map key type"),
302            "unexpected error: {msg}"
303        );
304    }
305}
risingwave_expr/aggregate/user_defined.rs

risingwave_expr/aggregate/
user_defined.rs
