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::{Array as _, 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        ensure_single_row(&arrow_output, "output")?;
89        let output = UdfArrowConvert::default().from_array(&self.return_field, &arrow_output)?;
90        // The UDF runtime is external input: a server may drift from the signature it was
91        // checked against at creation time. Surface a mistyped result instead of letting it
92        // corrupt downstream value encoding.
93        if output.data_type() != self.return_type {
94            return Err(anyhow::anyhow!(
95                "UDF returned a value of type {} while the declared return type is {}",
96                output.data_type(),
97                self.return_type
98            )
99            .into());
100        }
101        Ok(output.datum_at(0))
102    }
103
104    /// Encode the state into a datum that can be stored in state table.
105    fn encode_state(&self, state: &AggregateState) -> Result<Datum> {
106        let state = &state.downcast_ref::<State>().0;
107        ensure_single_row(state, "state")?;
108        let state = UdfArrowConvert::default().from_array(&self.state_field, state)?;
109        Ok(state.datum_at(0))
110    }
111
112    /// Decode the state from a datum in state table.
113    fn decode_state(&self, datum: Datum) -> Result<AggregateState> {
114        let array = {
115            let mut builder = DataType::Bytea.create_array_builder(1);
116            builder.append(datum);
117            builder.finish()
118        };
119        let state = UdfArrowConvert::default().to_array(self.state_field.data_type(), &array)?;
120        Ok(AggregateState::Any(Box::new(State(state))))
121    }
122}
123
124/// The runtime is external input: reject an array of the wrong length so that `datum_at(0)` at
125/// the call sites is in bounds.
126fn ensure_single_row(array: &ArrayRef, what: &str) -> Result<()> {
127    if array.len() != 1 {
128        return Err(anyhow::anyhow!(
129            "UDF aggregate {what} has {} rows, but expected exactly 1",
130            array.len()
131        )
132        .into());
133    }
134    Ok(())
135}
136
137// In arrow-udf, aggregate state is represented as an `ArrayRef`.
138// To avoid unnecessary conversion between `ArrayRef` and `Datum`,
139// we store `ArrayRef` directly in our `AggregateState`.
140#[derive(Debug)]
141struct State(ArrayRef);
142
143impl EstimateSize for State {
144    fn estimated_heap_size(&self) -> usize {
145        self.0.get_array_memory_size()
146    }
147}
148
149impl AggStateDyn for State {}
150
151/// Create a new user-defined aggregate function.
152pub fn new_user_defined(
153    return_type: &DataType,
154    udf: &PbUserDefinedFunctionMetadata,
155) -> Result<BoxedAggregateFunction> {
156    let arg_types = udf.arg_types.iter().map(|t| t.into()).collect::<Vec<_>>();
157    let language = udf.language.as_str();
158    let runtime = udf.runtime.as_deref();
159    let link = udf.link.as_deref();
160
161    let name_in_runtime = udf
162        .name_in_runtime()
163        .expect("SQL UDF won't get here, other UDFs must have `name_in_runtime`");
164
165    let build_fn = crate::sig::find_udf_impl(language, runtime, link)?.build_fn;
166    let runtime = build_fn(BuildOptions {
167        kind: UdfKind::Aggregate,
168        body: udf.body.as_deref(),
169        compressed_binary: udf.compressed_binary.as_deref(),
170        link: udf.link.as_deref(),
171        name_in_runtime,
172        arg_names: &udf.arg_names,
173        arg_types: &arg_types,
174        return_type,
175        always_retry_on_network_error: false,
176        language,
177        is_async: udf.is_async,
178        is_batched: udf.is_batched,
179    })
180    .context("failed to build UDF runtime")?;
181
182    // legacy UDF runtimes do not support aggregate functions,
183    // so we can assume that the runtime is not legacy
184    let arrow_convert = UdfArrowConvert::default();
185    let arg_schema = Arc::new(Schema::new(
186        arg_types
187            .iter()
188            .map(|t| arrow_convert.to_arrow_field("", t))
189            .try_collect::<_, Fields, _>()?,
190    ));
191
192    Ok(Box::new(UserDefinedAggregateFunction {
193        return_field: arrow_convert.to_arrow_field("", return_type)?,
194        state_field: Field::new(
195            "state",
196            risingwave_common::array::arrow::arrow_schema_udf::DataType::Binary,
197            true,
198        ),
199        return_type: return_type.clone(),
200        arg_schema,
201        runtime,
202    }))
203}
204
205#[cfg(test)]
206mod tests {
207    use anyhow::Result;
208    use futures::stream::BoxStream;
209    use risingwave_common::array::arrow::arrow_array_udf::{
210        BinaryArray, Float64Array, Int32Array, Int64Array, MapArray, RecordBatch, StringArray,
211        StructArray,
212    };
213    use risingwave_common::array::arrow::arrow_buffer_udf::OffsetBuffer;
214    use risingwave_common::array::arrow::arrow_schema_udf::DataType as ArrowDataType;
215    use risingwave_common::types::{MapType, StructType};
216
217    use super::*;
218    use crate::sig::UdfImpl;
219
220    /// Returns the given array from `finish`, regardless of the declared signature.
221    #[derive(Debug)]
222    struct StubRuntime(ArrayRef);
223
224    #[async_trait::async_trait]
225    impl UdfImpl for StubRuntime {
226        async fn call(&self, _input: &RecordBatch) -> Result<RecordBatch> {
227            unimplemented!()
228        }
229
230        async fn call_table_function<'a>(
231            &'a self,
232            _input: &'a RecordBatch,
233        ) -> Result<BoxStream<'a, Result<RecordBatch>>> {
234            unimplemented!()
235        }
236
237        async fn call_agg_finish(&self, _state: &ArrayRef) -> Result<ArrayRef> {
238            Ok(self.0.clone())
239        }
240    }
241
242    fn build_agg(return_type: DataType, finish_output: ArrayRef) -> UserDefinedAggregateFunction {
243        let convert = UdfArrowConvert::default();
244        UserDefinedAggregateFunction {
245            return_field: convert.to_arrow_field("", &return_type).unwrap(),
246            state_field: Field::new("state", ArrowDataType::Binary, true),
247            return_type,
248            arg_schema: Arc::new(Schema::new(Vec::<Field>::new())),
249            runtime: Box::new(StubRuntime(finish_output)),
250        }
251    }
252
253    /// Drives `get_result` with a runtime that returns `output` and yields the error message.
254    async fn get_result_err(return_type: DataType, output: ArrayRef) -> String {
255        let agg = build_agg(return_type, output);
256        let state = AggregateState::Any(Box::new(State(
257            Arc::new(Int64Array::from(vec![0i64])) as ArrayRef
258        )));
259        let err = agg.get_result(&state).await.unwrap_err();
260        format!("{:?}", anyhow::anyhow!(err))
261    }
262
263    /// Drives `encode_state` with `state` and yields the error message.
264    fn encode_state_err(state: ArrayRef) -> String {
265        let agg = build_agg(
266            DataType::Int64,
267            Arc::new(Int64Array::from(vec![Some(1i64)])),
268        );
269        let err = agg
270            .encode_state(&AggregateState::Any(Box::new(State(state))))
271            .unwrap_err();
272        format!("{:?}", anyhow::anyhow!(err))
273    }
274
275    #[tokio::test]
276    async fn misbehaving_runtime_mistyped_finish() {
277        // Struct child diverges from the declared `struct<a bigint>`.
278        let fields: Fields = vec![Field::new("a", ArrowDataType::Utf8, true)].into();
279        let struct_output: ArrayRef = Arc::new(StructArray::new(
280            fields,
281            vec![Arc::new(StringArray::from(vec![Some("oops")])) as ArrayRef],
282            None,
283        ));
284        let msg = get_result_err(
285            DataType::Struct(StructType::new(vec![("a", DataType::Int64)])),
286            struct_output,
287        )
288        .await;
289        assert!(
290            msg.contains("declared return type"),
291            "unexpected error: {msg}"
292        );
293
294        // Scalar output diverges from the declared `bigint`.
295        let msg = get_result_err(DataType::Int64, Arc::new(Int32Array::from(vec![Some(1)]))).await;
296        assert!(
297            msg.contains("declared return type"),
298            "unexpected error: {msg}"
299        );
300
301        // Map key type is unrepresentable in RW: must be a graceful error, not a panic
302        // from `MapArray::data_type()`.
303        let entries_fields: Fields = vec![
304            Field::new("key", ArrowDataType::Float64, false),
305            Field::new("value", ArrowDataType::Int32, true),
306        ]
307        .into();
308        let entries = StructArray::new(
309            entries_fields.clone(),
310            vec![
311                Arc::new(Float64Array::from(vec![Some(1.5)])) as ArrayRef,
312                Arc::new(Int32Array::from(vec![Some(42)])),
313            ],
314            None,
315        );
316        let map_output: ArrayRef = Arc::new(MapArray::new(
317            Arc::new(Field::new(
318                "entries",
319                ArrowDataType::Struct(entries_fields),
320                false,
321            )),
322            OffsetBuffer::new(vec![0, 1].into()),
323            entries,
324            None,
325            false,
326        ));
327        let msg = get_result_err(
328            DataType::Map(MapType::from_kv(DataType::Varchar, DataType::Int32)),
329            map_output,
330        )
331        .await;
332        assert!(
333            msg.contains("invalid map key type"),
334            "unexpected error: {msg}"
335        );
336    }
337
338    #[tokio::test]
339    async fn misbehaving_runtime_miscounted_finish() {
340        let msg = get_result_err(
341            DataType::Int64,
342            Arc::new(Int64Array::from(Vec::<Option<i64>>::new())),
343        )
344        .await;
345        assert!(msg.contains("0 rows"), "unexpected error: {msg}");
346
347        // More than one row does not panic, but `datum_at(0)` would silently drop the rest.
348        let msg = get_result_err(
349            DataType::Int64,
350            Arc::new(Int64Array::from(vec![Some(1i64), Some(2i64)])),
351        )
352        .await;
353        assert!(msg.contains("2 rows"), "unexpected error: {msg}");
354    }
355
356    #[test]
357    fn misbehaving_runtime_miscounted_state() {
358        let msg = encode_state_err(Arc::new(BinaryArray::from(Vec::<Option<&[u8]>>::new())));
359        assert!(msg.contains("0 rows"), "unexpected error: {msg}");
360
361        let msg = encode_state_err(Arc::new(BinaryArray::from(vec![
362            Some(b"a".as_slice()),
363            Some(b"b".as_slice()),
364        ])));
365        assert!(msg.contains("2 rows"), "unexpected error: {msg}");
366    }
367}
risingwave_expr/aggregate/user_defined.rs

risingwave_expr/aggregate/
user_defined.rs
