risingwave_expr_impl/scalar/
array_access.rs

1// Copyright 2025 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 risingwave_common::array::ListRef;
16use risingwave_common::types::ScalarRefImpl;
17use risingwave_expr::function;
18
19#[function("array_access(anyarray, int4) -> any")]
20fn array_access(list: ListRef<'_>, index: i32) -> Option<ScalarRefImpl<'_>> {
21    // index must be greater than 0 following a one-based numbering convention for arrays
22    if index < 1 {
23        return None;
24    }
25    // returns `NULL` if index is out of bounds
26    list.get(index as usize - 1).flatten()
27}
28
29#[cfg(test)]
30mod tests {
31
32    use risingwave_common::array::ListValue;
33    use risingwave_common::types::Scalar;
34
35    use super::*;
36
37    #[test]
38    fn test_int4_array_access() {
39        let v1 = ListValue::from_iter([1, 2, 3]);
40        let l1 = v1.as_scalar_ref();
41
42        assert_eq!(array_access(l1, 1), Some(1.into()));
43        assert_eq!(array_access(l1, -1), None);
44        assert_eq!(array_access(l1, 0), None);
45        assert_eq!(array_access(l1, 4), None);
46    }
47
48    #[test]
49    fn test_utf8_array_access() {
50        let v1 = ListValue::from_iter(["来自", "foo", "bar"]);
51        let v2 = ListValue::from_iter(["fizz", "荷兰", "buzz"]);
52        let v3 = ListValue::from_iter([None, None, Some("的爱")]);
53
54        assert_eq!(array_access(v1.as_scalar_ref(), 1), Some("来自".into()));
55        assert_eq!(array_access(v2.as_scalar_ref(), 2), Some("荷兰".into()));
56        assert_eq!(array_access(v3.as_scalar_ref(), 3), Some("的爱".into()));
57    }
58
59    #[test]
60    fn test_nested_array_access() {
61        let v = ListValue::from_iter([
62            ListValue::from_iter(["foo", "bar"]),
63            ListValue::from_iter(["fizz", "buzz"]),
64        ]);
65        assert_eq!(
66            array_access(v.as_scalar_ref(), 1),
67            Some(ListValue::from_iter(["foo", "bar"]).as_scalar_ref().into())
68        );
69    }
70}
risingwave_expr_impl/scalar/array_access.rs

risingwave_expr_impl/scalar/
array_access.rs
