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;
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(
68    array: ListRef<'_>,
69    ctx: &Context,
70    writer: &mut impl risingwave_common::array::ListWrite,
71) -> Result<()> {
72    // The elements of the array must be arrays themselves
73    let outer_type = &ctx.arg_types[0];
74    let inner_type = if outer_type.is_array() {
75        outer_type.as_list_elem()
76    } else {
77        return Err(ExprError::InvalidParam {
78            name: "array_flatten",
79            reason: Box::from("expected the argument to be an array of arrays"),
80        });
81    };
82    if !inner_type.is_array() {
83        return Err(ExprError::InvalidParam {
84            name: "array_flatten",
85            reason: Box::from("expected the argument to be an array of arrays"),
86        });
87    }
88
89    writer.write_iter(
90        array
91            .iter()
92            // Filter out NULL inner arrays
93            .flatten()
94            // Flatten all inner arrays
95            .flat_map(|inner_array| inner_array.into_list().iter()),
96    );
97    Ok(())
98}
risingwave_expr_impl/scalar/array_flatten.rs

risingwave_expr_impl/scalar/
array_flatten.rs
