risingwave_common/array/arrow/
arrow_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//! This is for arrow dependency named `arrow-xxx` such as `arrow-array` in the cargo workspace.
16//!
17//! This should the default arrow version to be used in our system.
18//!
19//! The corresponding version of arrow is currently used by `udf`.
20
21use std::sync::Arc;
22
23pub use super::arrow_58::{
24    FromArrow, ToArrow, arrow_array, arrow_buffer, arrow_cast, arrow_schema,
25};
26use crate::array::{ArrayError, ArrayImpl, DataType, DecimalArray, JsonbArray};
27
28/// Arrow conversion for UDF.
29#[derive(Default, Debug)]
30pub struct UdfArrowConvert {
31    /// Whether the UDF talks in legacy mode.
32    ///
33    /// If true, decimal and jsonb types are mapped to Arrow `LargeBinary` and `LargeUtf8` types.
34    /// Otherwise, they are mapped to Arrow extension types.
35    /// See <https://github.com/risingwavelabs/arrow-udf/tree/main#extension-types>.
36    pub legacy: bool,
37}
38
39impl ToArrow for UdfArrowConvert {
40    fn decimal_to_arrow(
41        &self,
42        _data_type: &arrow_schema::DataType,
43        array: &DecimalArray,
44    ) -> Result<arrow_array::ArrayRef, ArrayError> {
45        if self.legacy {
46            // Decimal values are stored as ASCII text representation in a large binary array.
47            Ok(Arc::new(arrow_array::LargeBinaryArray::from(array)))
48        } else {
49            Ok(Arc::new(arrow_array::StringArray::from(array)))
50        }
51    }
52
53    fn jsonb_to_arrow(&self, array: &JsonbArray) -> Result<arrow_array::ArrayRef, ArrayError> {
54        if self.legacy {
55            // JSON values are stored as text representation in a large string array.
56            Ok(Arc::new(arrow_array::LargeStringArray::from(array)))
57        } else {
58            Ok(Arc::new(arrow_array::StringArray::from(array)))
59        }
60    }
61
62    fn jsonb_type_to_arrow(&self, name: &str) -> arrow_schema::Field {
63        if self.legacy {
64            arrow_schema::Field::new(name, arrow_schema::DataType::LargeUtf8, true)
65        } else {
66            arrow_schema::Field::new(name, arrow_schema::DataType::Utf8, true)
67                .with_metadata([("ARROW:extension:name".into(), "arrowudf.json".into())].into())
68        }
69    }
70
71    fn decimal_type_to_arrow(&self, name: &str) -> arrow_schema::Field {
72        if self.legacy {
73            arrow_schema::Field::new(name, arrow_schema::DataType::LargeBinary, true)
74        } else {
75            arrow_schema::Field::new(name, arrow_schema::DataType::Utf8, true)
76                .with_metadata([("ARROW:extension:name".into(), "arrowudf.decimal".into())].into())
77        }
78    }
79}
80
81impl FromArrow for UdfArrowConvert {
82    fn from_large_utf8(&self) -> Result<DataType, ArrayError> {
83        if self.legacy {
84            Ok(DataType::Jsonb)
85        } else {
86            Ok(DataType::Varchar)
87        }
88    }
89
90    fn from_large_binary(&self) -> Result<DataType, ArrayError> {
91        if self.legacy {
92            Ok(DataType::Decimal)
93        } else {
94            Ok(DataType::Bytea)
95        }
96    }
97
98    fn from_large_utf8_array(
99        &self,
100        array: &arrow_array::LargeStringArray,
101    ) -> Result<ArrayImpl, ArrayError> {
102        if self.legacy {
103            Ok(ArrayImpl::Jsonb(array.try_into()?))
104        } else {
105            Ok(ArrayImpl::Utf8(array.into()))
106        }
107    }
108
109    fn from_large_binary_array(
110        &self,
111        array: &arrow_array::LargeBinaryArray,
112    ) -> Result<ArrayImpl, ArrayError> {
113        if self.legacy {
114            Ok(ArrayImpl::Decimal(array.try_into()?))
115        } else {
116            Ok(ArrayImpl::Bytea(array.into()))
117        }
118    }
119}
120
121#[cfg(test)]
122mod tests {
123
124    use arrow_array::Array as _;
125
126    use super::*;
127    use crate::array::*;
128
129    /// Wraps an array's own type into a nameless field for the `from_*_array` calls.
130    fn typed_field(array: &impl arrow_array::Array) -> arrow_schema::Field {
131        arrow_schema::Field::new("", array.data_type().clone(), true)
132    }
133
134    #[test]
135    fn struct_array() {
136        // Empty array - risingwave to arrow conversion.
137        let test_arr = StructArray::new(StructType::empty(), vec![], Bitmap::ones(0));
138        assert_eq!(
139            UdfArrowConvert::default()
140                .struct_to_arrow(
141                    &arrow_schema::DataType::Struct(arrow_schema::Fields::empty()),
142                    &test_arr
143                )
144                .unwrap()
145                .len(),
146            0
147        );
148
149        // Empty array - arrow to risingwave conversion.
150        let test_arr_2 = arrow_array::StructArray::new_empty_fields(0, None);
151        let test_arr_2_field = typed_field(&test_arr_2);
152        assert_eq!(
153            UdfArrowConvert::default()
154                .from_struct_array(&test_arr_2_field, &test_arr_2)
155                .unwrap()
156                .len(),
157            0
158        );
159
160        // Struct array with primitive types. arrow to risingwave conversion.
161        let test_arrow_struct_array = arrow_array::StructArray::try_from(vec![
162            (
163                "a",
164                Arc::new(arrow_array::BooleanArray::from(vec![
165                    Some(false),
166                    Some(false),
167                    Some(true),
168                    None,
169                ])) as arrow_array::ArrayRef,
170            ),
171            (
172                "b",
173                Arc::new(arrow_array::Int32Array::from(vec![
174                    Some(42),
175                    Some(28),
176                    Some(19),
177                    None,
178                ])) as arrow_array::ArrayRef,
179            ),
180        ])
181        .unwrap();
182        let struct_field = typed_field(&test_arrow_struct_array);
183        let actual_risingwave_struct_array = UdfArrowConvert::default()
184            .from_struct_array(&struct_field, &test_arrow_struct_array)
185            .unwrap()
186            .into_struct();
187        let expected_risingwave_struct_array = StructArray::new(
188            StructType::new(vec![("a", DataType::Boolean), ("b", DataType::Int32)]),
189            vec![
190                BoolArray::from_iter([Some(false), Some(false), Some(true), None]).into_ref(),
191                I32Array::from_iter([Some(42), Some(28), Some(19), None]).into_ref(),
192            ],
193            [true, true, true, true].into_iter().collect(),
194        );
195        assert_eq!(
196            expected_risingwave_struct_array,
197            actual_risingwave_struct_array
198        );
199    }
200
201    #[test]
202    fn list() {
203        let array = ListArray::from_iter([None, Some(vec![0, -127, 127, 50]), Some(vec![0; 0])]);
204        let data_type = arrow_schema::DataType::new_list(arrow_schema::DataType::Int32, true);
205        let arrow = UdfArrowConvert::default()
206            .list_to_arrow(&data_type, &array)
207            .unwrap();
208        let list_field = typed_field(&arrow);
209        let rw_array = UdfArrowConvert::default()
210            .from_list_array(&list_field, arrow.as_any().downcast_ref().unwrap())
211            .unwrap();
212        assert_eq!(rw_array.as_list(), &array);
213    }
214
215    #[test]
216    fn map() {
217        let map_type = MapType::from_kv(DataType::Varchar, DataType::Int32);
218        let rw_map_type = DataType::Map(map_type.clone());
219        let mut builder = MapArrayBuilder::with_type(3, rw_map_type.clone());
220        builder.append_owned(Some(
221            MapValue::try_from_kv(
222                ListValue::from_str("{a,b,c}", &DataType::Varchar.list()).unwrap(),
223                ListValue::from_str("{1,2,3}", &DataType::Int32.list()).unwrap(),
224            )
225            .unwrap(),
226        ));
227        builder.append_owned(None);
228        builder.append_owned(Some(
229            MapValue::try_from_kv(
230                ListValue::from_str("{a,c}", &DataType::Varchar.list()).unwrap(),
231                ListValue::from_str("{1,3}", &DataType::Int32.list()).unwrap(),
232            )
233            .unwrap(),
234        ));
235        let rw_array = builder.finish();
236
237        let arrow_map_type = UdfArrowConvert::default()
238            .map_type_to_arrow(&map_type)
239            .unwrap();
240        expect_test::expect![[r#"
241            Map(
242                Field {
243                    name: "entries",
244                    data_type: Struct(
245                        [
246                            Field {
247                                name: "key",
248                                data_type: Utf8,
249                            },
250                            Field {
251                                name: "value",
252                                data_type: Int32,
253                                nullable: true,
254                            },
255                        ],
256                    ),
257                },
258                false,
259            )
260        "#]]
261        .assert_debug_eq(&arrow_map_type);
262        let rw_map_type_new = UdfArrowConvert::default()
263            .from_field(&arrow_schema::Field::new(
264                "map",
265                arrow_map_type.clone(),
266                true,
267            ))
268            .unwrap();
269        assert_eq!(rw_map_type, rw_map_type_new);
270        let arrow = UdfArrowConvert::default()
271            .map_to_arrow(&arrow_map_type, &rw_array)
272            .unwrap();
273        expect_test::expect![[r#"
274            MapArray
275            [
276              StructArray
277            -- validity:
278            [
279              valid,
280              valid,
281              valid,
282            ]
283            [
284            -- child 0: "key" (Utf8)
285            StringArray
286            [
287              "a",
288              "b",
289              "c",
290            ]
291            -- child 1: "value" (Int32)
292            PrimitiveArray<Int32>
293            [
294              1,
295              2,
296              3,
297            ]
298            ],
299              null,
300              StructArray
301            -- validity:
302            [
303              valid,
304              valid,
305            ]
306            [
307            -- child 0: "key" (Utf8)
308            StringArray
309            [
310              "a",
311              "c",
312            ]
313            -- child 1: "value" (Int32)
314            PrimitiveArray<Int32>
315            [
316              1,
317              3,
318            ]
319            ],
320            ]"#]]
321        .assert_eq(
322            &format!("{:#?}", arrow)
323                .lines()
324                .map(|s| s.trim_end())
325                .collect::<Vec<_>>()
326                .join("\n"),
327        );
328
329        let map_field = typed_field(&arrow);
330        let rw_array_new = UdfArrowConvert::default()
331            .from_map_array(&map_field, arrow.as_any().downcast_ref().unwrap())
332            .unwrap();
333        assert_eq!(&rw_array, rw_array_new.as_map());
334    }
335}
risingwave_common/array/arrow/arrow_udf.rs

risingwave_common/array/arrow/
arrow_udf.rs
