risingwave_expr_impl/scalar/
array_flatten.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, ListValue};
16use risingwave_expr::expr::Context;
17use risingwave_expr::{ExprError, Result, function};
18
19/// Flattens a nested array by concatenating the inner arrays into a single array.
20/// Only the outermost level of nesting is removed. For deeper nested arrays, call
21/// `array_flatten` multiple times.
22///
23/// Examples:
24///
25/// ```slt
26/// query T
27/// select array_flatten(array[array[1, 2], array[3, 4]]);
28/// ----
29/// {1,2,3,4}
30///
31/// query T
32/// select array_flatten(array[array[1, 2], array[]::int[], array[3, 4]]);
33/// ----
34/// {1,2,3,4}
35///
36/// query T
37/// select array_flatten(array[array[1, 2], null, array[3, 4]]);
38/// ----
39/// {1,2,3,4}
40///
41/// query T
42/// select array_flatten(array[array[array[1], array[2, null]], array[array[3, 4], null::int[]]]);
43/// ----
44/// {{1},{2,NULL},{3,4},NULL}
45///
46/// query T
47/// select array_flatten(array[[]]::int[][]);
48/// ----
49/// {}
50///
51/// query T
52/// select array_flatten(array[[null, 1]]::int[][]);
53/// ----
54/// {NULL,1}
55///
56/// query T
57/// select array_flatten(array[]::int[][]);
58/// ----
59/// {}
60///
61/// query T
62/// select array_flatten(null::int[][]);
63/// ----
64/// NULL
65/// ```
66#[function("array_flatten(anyarray) -> anyarray")]
67fn array_flatten(array: ListRef<'_>, ctx: &Context) -> Result<ListValue> {
68    // The elements of the array must be arrays themselves
69    let outer_type = &ctx.arg_types[0];
70    let inner_type = if outer_type.is_array() {
71        outer_type.as_list_element_type()
72    } else {
73        return Err(ExprError::InvalidParam {
74            name: "array_flatten",
75            reason: Box::from("expected the argument to be an array of arrays"),
76        });
77    };
78    if !inner_type.is_array() {
79        return Err(ExprError::InvalidParam {
80            name: "array_flatten",
81            reason: Box::from("expected the argument to be an array of arrays"),
82        });
83    }
84    let inner_elem_type = inner_type.as_list_element_type();
85
86    // Collect all inner array elements and flatten them into a single array
87    Ok(ListValue::from_datum_iter(
88        inner_elem_type,
89        array
90            .iter()
91            // Filter out NULL inner arrays
92            .flatten()
93            // Flatten all inner arrays
94            .flat_map(|inner_array| inner_array.into_list().iter()),
95    ))
96}
risingwave_expr_impl/scalar/array_flatten.rs

risingwave_expr_impl/scalar/
array_flatten.rs
