risingwave_expr_impl/scalar/
array_concat.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::expr::Context;
18use risingwave_expr::function;
19
20/// Concatenates the two arrays.
21///
22/// Examples:
23///
24/// ```slt
25/// # concat
26/// query T
27/// select array_cat(array[66], array[123]);
28/// ----
29/// {66,123}
30///
31/// query T
32/// select array_cat(array[66], null::int[]);
33/// ----
34/// {66}
35///
36/// query T
37/// select array_cat(null::int[], array[123]);
38/// ----
39/// {123}
40///
41/// query T
42/// select array_cat(null::int[], null::int[]);
43/// ----
44/// NULL
45///
46/// # append
47/// query T
48/// select array_cat(array[array[66]], array[233]);
49/// ----
50/// {{66},{233}}
51///
52/// query T
53/// select array_cat(array[array[66]], null::int[]);
54/// ----
55/// {{66}}
56///
57/// # different from PG
58/// query T
59/// select array_cat(null::int[][], array[233]);
60/// ----
61/// {{233}}
62///
63/// query T
64/// select array_cat(null::int[][], null::int[]);
65/// ----
66/// NULL
67///
68/// # prepend
69/// query T
70/// select array_cat(array[233], array[array[66]]);
71/// ----
72/// {{233},{66}}
73///
74/// query T
75/// select array_cat(null::int[], array[array[66]]);
76/// ----
77/// {{66}}
78///
79/// # different from PG
80/// query T
81/// select array_cat(array[233], null::int[][]);
82/// ----
83/// {{233}}
84///
85/// query T
86/// select array_cat(null::int[], null::int[][]);
87/// ----
88/// NULL
89/// ```
90#[function("array_cat(anyarray, anyarray) -> anyarray")]
91fn array_cat(
92    left: Option<ListRef<'_>>,
93    right: Option<ListRef<'_>>,
94    ctx: &Context,
95    writer: &mut impl risingwave_common::array::ListWrite,
96) -> Option<()> {
97    if ctx.arg_types[0] == ctx.arg_types[1] {
98        // array || array
99        if left.is_none() && right.is_none() {
100            return None;
101        }
102
103        if let Some(left) = left {
104            writer.write_iter(left.iter());
105        }
106        if let Some(right) = right {
107            writer.write_iter(right.iter());
108        }
109    } else if ctx.arg_types[0].as_list_elem() == &ctx.arg_types[1] {
110        // array[] || array
111        if left.is_none() && right.is_none() {
112            return None;
113        }
114
115        if let Some(left) = left {
116            writer.write_iter(left.iter());
117        }
118        if let Some(right) = right {
119            writer.write(Some(ScalarRefImpl::from(right)));
120        }
121    } else if &ctx.arg_types[0] == ctx.arg_types[1].as_list_elem() {
122        // array || array[]
123        if left.is_none() && right.is_none() {
124            return None;
125        }
126
127        if let Some(left) = left {
128            writer.write(Some(ScalarRefImpl::from(left)));
129        }
130        if let Some(right) = right {
131            writer.write_iter(right.iter());
132        }
133    } else {
134        unreachable!()
135    }
136    Some(())
137}
138
139/// Appends a value as the back element of an array.
140/// The behavior is the same as PG.
141///
142/// Examples:
143///
144/// ```slt
145/// query T
146/// select array_append(array[66], 123);
147/// ----
148/// {66,123}
149///
150/// query T
151/// select array_append(array[66], null::int);
152/// ----
153/// {66,NULL}
154///
155/// query T
156/// select array_append(null::int[], 233);
157/// ----
158/// {233}
159///
160/// query T
161/// select array_append(null::int[], null::int);
162/// ----
163/// {NULL}
164/// ```
165#[function("array_append(anyarray, any) -> anyarray")]
166fn array_append(
167    left: Option<ListRef<'_>>,
168    right: Option<ScalarRefImpl<'_>>,
169    writer: &mut impl risingwave_common::array::ListWrite,
170) {
171    writer.write_iter(left.iter().flat_map(|list| list.iter()));
172    writer.write(right);
173}
174
175/// Prepends a value as the front element of an array.
176/// The behavior is the same as PG.
177///
178/// Examples:
179///
180/// ```slt
181/// query T
182/// select array_prepend(123, array[66]);
183/// ----
184/// {123,66}
185///
186/// query T
187/// select array_prepend(null::int, array[66]);
188/// ----
189/// {NULL,66}
190///
191/// query T
192/// select array_prepend(233, null::int[]);
193/// ----
194/// {233}
195///
196/// query T
197/// select array_prepend(null::int, null::int[]);
198/// ----
199/// {NULL}
200/// ```
201#[function("array_prepend(any, anyarray) -> anyarray")]
202fn array_prepend(
203    left: Option<ScalarRefImpl<'_>>,
204    right: Option<ListRef<'_>>,
205    writer: &mut impl risingwave_common::array::ListWrite,
206) {
207    writer.write(left);
208    writer.write_iter(right.iter().flat_map(|list| list.iter()));
209}
risingwave_expr_impl/scalar/array_concat.rs

risingwave_expr_impl/scalar/
array_concat.rs
