feat(expr): implement array_flatten (#21640)

Signed-off-by: Richard Chien <[email protected]>

Author: stdrc

proto/expr.proto

@@ -236,6 +236,7 @@ message ExprNode {
     ARRAY_SORT = 549;
     ARRAY_CONTAINS = 550;
     ARRAY_CONTAINED = 551;
+    ARRAY_FLATTEN = 552;
 
     // Int256 functions
     HEX_TO_INT256 = 560;

src/common/src/types/mod.rs

@@ -510,13 +510,21 @@ impl DataType {
     /// # Panics
     ///
     /// Panics if the type is not a list type.
+    /// TODO(rc): rename to `as_list_element_type`
     pub fn as_list(&self) -> &DataType {
         match self {
             DataType::List(t) => t,
             t => panic!("expect list type, got {t}"),
         }
     }
 
+    pub fn into_list_element_type(self) -> DataType {
+        match self {
+            DataType::List(t) => *t,
+            t => panic!("expect list type, got {t}"),
+        }
+    }
+
     /// Return a new type that removes the outer list, and get the innermost element type.
     ///
     /// Use [`DataType::as_list`] if you only want the element type of a list.

src/expr/impl/src/scalar/array_flatten.rs

@@ -0,0 +1,96 @@
+// Copyright 2025 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use risingwave_common::array::{ListRef, ListValue};
+use risingwave_expr::expr::Context;
+use risingwave_expr::{ExprError, Result, function};
+
+/// Flattens a nested array by concatenating the inner arrays into a single array.
+/// Only the outermost level of nesting is removed. For deeper nested arrays, call
+/// `array_flatten` multiple times.
+///
+/// Examples:
+///
+/// ```slt
+/// query T
+/// select array_flatten(array[array[1, 2], array[3, 4]]);
+/// ----
+/// {1,2,3,4}
+///
+/// query T
+/// select array_flatten(array[array[1, 2], array[]::int[], array[3, 4]]);
+/// ----
+/// {1,2,3,4}
+///
+/// query T
+/// select array_flatten(array[array[1, 2], null, array[3, 4]]);
+/// ----
+/// {1,2,3,4}
+///
+/// query T
+/// select array_flatten(array[array[array[1], array[2, null]], array[array[3, 4], null::int[]]]);
+/// ----
+/// {{1},{2,NULL},{3,4},NULL}
+///
+/// query T
+/// select array_flatten(array[[]]::int[][]);
+/// ----
+/// {}
+///
+/// query T
+/// select array_flatten(array[[null, 1]]::int[][]);
+/// ----
+/// {NULL,1}
+///
+/// query T
+/// select array_flatten(array[]::int[][]);
+/// ----
+/// {}
+///
+/// query T
+/// select array_flatten(null::int[][]);
+/// ----
+/// NULL
+/// ```
+#[function("array_flatten(anyarray) -> anyarray")]
+fn array_flatten(array: ListRef<'_>, ctx: &Context) -> Result<ListValue> {
+    // The elements of the array must be arrays themselves
+    let outer_type = &ctx.arg_types[0];
+    let inner_type = if outer_type.is_array() {
+        outer_type.as_list()
+    } else {
+        return Err(ExprError::InvalidParam {
+            name: "array_flatten",
+            reason: Box::from("expected the argument to be an array of arrays"),
+        });
+    };
+    if !inner_type.is_array() {
+        return Err(ExprError::InvalidParam {
+            name: "array_flatten",
+            reason: Box::from("expected the argument to be an array of arrays"),
+        });
+    }
+    let inner_elem_type = inner_type.as_list();
+
+    // Collect all inner array elements and flatten them into a single array
+    Ok(ListValue::from_datum_iter(
+        inner_elem_type,
+        array
+            .iter()
+            // Filter out NULL inner arrays
+            .flatten()
+            // Flatten all inner arrays
+            .flat_map(|inner_array| inner_array.into_list().iter()),
+    ))
+}

src/expr/impl/src/scalar/mod.rs

@@ -18,6 +18,7 @@ mod array_access;
 mod array_concat;
 mod array_contain;
 mod array_distinct;
+mod array_flatten;
 mod array_length;
 mod array_min_max;
 mod array_positions;

src/frontend/src/binder/expr/function/builtin_scalar.rs

@@ -344,6 +344,15 @@ impl Binder {
                 ("arraycontains", raw_call(ExprType::ArrayContains)),
                 ("array_contained", raw_call(ExprType::ArrayContained)),
                 ("arraycontained", raw_call(ExprType::ArrayContained)),
+                ("array_flatten", guard_by_len(1, raw(|_binder, inputs| {
+                    inputs[0].ensure_array_type().map_err(|_| ErrorCode::BindError("array_flatten expects `any[][]` input".into()))?;
+                    let return_type = inputs[0].return_type().into_list_element_type();
+                    if !return_type.is_array() {
+                        return Err(ErrorCode::BindError("array_flatten expects `any[][]` input".into()).into());
+
+                    }
+                    Ok(FunctionCall::new_unchecked(ExprType::ArrayFlatten, inputs, return_type).into())
+                }))),
                 ("trim_array", raw_call(ExprType::TrimArray)),
                 (
                     "array_ndims",

src/frontend/src/expr/pure.rs

@@ -183,6 +183,7 @@ impl ExprVisitor for ImpureAnalyzer {
             | Type::ArrayPosition
             | Type::ArrayContains
             | Type::ArrayContained
+            | Type::ArrayFlatten
             | Type::HexToInt256
             | Type::JsonbConcat
             | Type::JsonbAccess

src/frontend/src/optimizer/plan_expr_visitor/strong.rs

@@ -264,6 +264,7 @@ impl Strong {
             | ExprType::ArraySort
             | ExprType::ArrayContains
             | ExprType::ArrayContained
+            | ExprType::ArrayFlatten
             | ExprType::HexToInt256
             | ExprType::JsonbAccess
             | ExprType::JsonbAccessStr