feat(expr): support VARIADIC function arguments

proto/expr.proto

@@ -91,6 +91,7 @@ message ExprNode {
     TRANSLATE = 216;
     COALESCE = 217;
     CONCAT_WS = 218;
+    CONCAT_WS_VARIADIC = 285;
     ABS = 219;
     SPLIT_PART = 220;
     CEIL = 221;
@@ -100,6 +101,8 @@ message ExprNode {
     CHAR_LENGTH = 225;
     REPEAT = 226;
     CONCAT_OP = 227;
+    CONCAT = 286;
+    CONCAT_VARIADIC = 287;
     // BOOL_OUT is different from CAST-bool-to-varchar in PostgreSQL.
     BOOL_OUT = 228;
     OCTET_LENGTH = 229;
@@ -176,6 +179,7 @@ message ExprNode {
     LEFT = 317;
     RIGHT = 318;
     FORMAT = 319;
+    FORMAT_VARIADIC = 324;
     PGWIRE_SEND = 320;
     PGWIRE_RECV = 321;
     CONVERT_FROM = 322;
@@ -224,9 +228,11 @@ message ExprNode {
     // jsonb ->> int, jsonb ->> text that returns text
     JSONB_ACCESS_STR = 601;
     // jsonb #> text[] -> jsonb
-    JSONB_EXTRACT_PATH = 613;
+    JSONB_EXTRACT_PATH = 626;
+    JSONB_EXTRACT_PATH_VARIADIC = 613;
     // jsonb #>> text[] -> text
-    JSONB_EXTRACT_PATH_TEXT = 614;
+    JSONB_EXTRACT_PATH_TEXT = 627;
+    JSONB_EXTRACT_PATH_TEXT_VARIADIC = 614;
     JSONB_TYPEOF = 602;
     JSONB_ARRAY_LENGTH = 603;
     IS_JSON = 604;
@@ -253,7 +259,9 @@ message ExprNode {
     JSONB_STRIP_NULLS = 616;
     TO_JSONB = 617;
     JSONB_BUILD_ARRAY = 618;
+    JSONB_BUILD_ARRAY_VARIADIC = 624;
     JSONB_BUILD_OBJECT = 619;
+    JSONB_BUILD_OBJECT_VARIADIC = 625;
     JSONB_PATH_EXISTS = 620;
     JSONB_PATH_MATCH = 621;
     JSONB_PATH_QUERY_ARRAY = 622;

src/common/src/array/list_array.rs

@@ -612,6 +612,24 @@ impl Debug for ListRef<'_> {
     }
 }
 
+impl Row for ListRef<'_> {
+    fn datum_at(&self, index: usize) -> DatumRef<'_> {
+        self.array.value_at(self.start as usize + index)
+    }
+
+    unsafe fn datum_at_unchecked(&self, index: usize) -> DatumRef<'_> {
+        self.array.value_at_unchecked(self.start as usize + index)
+    }
+
+    fn len(&self) -> usize {
+        self.len()
+    }
+
+    fn iter(&self) -> impl Iterator<Item = DatumRef<'_>> {
+        (*self).iter()
+    }
+}
+
 impl ToText for ListRef<'_> {
     // This function will be invoked when pgwire prints a list value in string.
     // Refer to PostgreSQL `array_out` or `appendPGArray`.

src/expr/core/src/expr/mod.rs

@@ -209,4 +209,6 @@ where
 pub struct Context {
     pub arg_types: Vec<DataType>,
     pub return_type: DataType,
+    /// Whether the function is variadic.
+    pub variadic: bool,
 }

src/expr/impl/src/scalar/cast.rs

@@ -301,6 +301,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("int[]").unwrap(),
+            variadic: false,
         };
         assert_eq!(
             str_to_list("{}", &ctx).unwrap(),
@@ -313,6 +314,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("int[]").unwrap(),
+            variadic: false,
         };
         assert_eq!(str_to_list("{1, 2, 3}", &ctx).unwrap(), list123);
 
@@ -321,6 +323,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("int[][]").unwrap(),
+            variadic: false,
         };
         assert_eq!(str_to_list("{{1, 2, 3}}", &ctx).unwrap(), nested_list123);
 
@@ -333,6 +336,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("int[][][]").unwrap(),
+            variadic: false,
         };
         assert_eq!(
             str_to_list("{{{1, 2, 3}}, {{44, 55, 66}}}", &ctx).unwrap(),
@@ -343,6 +347,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::from_str("int[][]").unwrap()],
             return_type: DataType::from_str("varchar[][]").unwrap(),
+            variadic: false,
         };
         let double_nested_varchar_list123_445566 = ListValue::from_iter([
             list_cast(nested_list123.as_scalar_ref(), &ctx).unwrap(),
@@ -353,6 +358,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("varchar[][][]").unwrap(),
+            variadic: false,
         };
         assert_eq!(
             str_to_list("{{{1, 2, 3}}, {{44, 55, 66}}}", &ctx).unwrap(),
@@ -366,6 +372,7 @@ mod tests {
         let ctx = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::from_str("int[]").unwrap(),
+            variadic: false,
         };
         assert!(str_to_list("{{}", &ctx).is_err());
         assert!(str_to_list("{}}", &ctx).is_err());
@@ -384,6 +391,7 @@ mod tests {
                 ("a", DataType::Int32),
                 ("b", DataType::Int32),
             ])),
+            variadic: false,
         };
         assert_eq!(
             struct_cast(
@@ -419,6 +427,7 @@ mod tests {
         let ctx_str_to_int16 = Context {
             arg_types: vec![DataType::Varchar],
             return_type: DataType::Int16,
+            variadic: false,
         };
         test_str_to_int16::<I16Array, _>(|x| str_parse(x, &ctx_str_to_int16).unwrap()).await;
     }

src/expr/impl/src/scalar/concat.rs

@@ -0,0 +1,73 @@
+// Copyright 2024 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 std::fmt::Write;
+
+use risingwave_common::row::Row;
+use risingwave_common::types::ToText;
+use risingwave_expr::function;
+
+/// Concatenates the text representations of all the arguments. NULL arguments are ignored.
+///
+/// # Example
+///
+/// ```slt
+/// query T
+/// select concat('abcde', 2, NULL, 22);
+/// ----
+/// abcde222
+///
+/// query T
+/// select concat(variadic array['abcde', '2', NULL, '22']);
+/// ----
+/// abcde222
+/// ```
+#[function("concat(variadic anyarray) -> varchar")]
+fn concat(vals: impl Row, writer: &mut impl Write) {
+    for string in vals.iter().flatten() {
+        string.write(writer).unwrap();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use risingwave_common::array::DataChunk;
+    use risingwave_common::row::Row;
+    use risingwave_common::test_prelude::DataChunkTestExt;
+    use risingwave_common::types::ToOwnedDatum;
+    use risingwave_common::util::iter_util::ZipEqDebug;
+    use risingwave_expr::expr::build_from_pretty;
+
+    #[tokio::test]
+    async fn test_concat() {
+        let concat = build_from_pretty("(concat:varchar $0:varchar $1:varchar $2:varchar)");
+        let (input, expected) = DataChunk::from_pretty(
+            "T T T  T
+             a b c  abc
+             . b c  bc
+             . . .  (empty)",
+        )
+        .split_column_at(3);
+
+        // test eval
+        let output = concat.eval(&input).await.unwrap();
+        assert_eq!(&output, expected.column_at(0));
+
+        // test eval_row
+        for (row, expected) in input.rows().zip_eq_debug(expected.rows()) {
+            let result = concat.eval_row(&row.to_owned_row()).await.unwrap();
+            assert_eq!(result, expected.datum_at(0).to_owned_datum());
+        }
+    }
+}

src/expr/impl/src/scalar/concat_ws.rs

@@ -20,8 +20,22 @@ use risingwave_expr::function;
 
 /// Concatenates all but the first argument, with separators. The first argument is used as the
 /// separator string, and should not be NULL. Other NULL arguments are ignored.
-#[function("concat_ws(varchar, ...) -> varchar")]
-fn concat_ws(sep: &str, vals: impl Row, writer: &mut impl Write) -> Option<()> {
+///
+/// # Example
+///
+/// ```slt
+/// query T
+/// select concat_ws(',', 'abcde', 2, NULL, 22);
+/// ----
+/// abcde,2,22
+///
+/// query T
+/// select concat_ws(',', variadic array['abcde', 2, NULL, 22] :: varchar[]);
+/// ----
+/// abcde,2,22
+/// ```
+#[function("concat_ws(varchar, variadic anyarray) -> varchar")]
+fn concat_ws(sep: &str, vals: impl Row, writer: &mut impl Write) {
     let mut string_iter = vals.iter().flatten();
     if let Some(string) = string_iter.next() {
         string.write(writer).unwrap();
@@ -30,7 +44,6 @@ fn concat_ws(sep: &str, vals: impl Row, writer: &mut impl Write) -> Option<()> {
         write!(writer, "{}", sep).unwrap();
         string.write(writer).unwrap();
     }
-    Some(())
 }
 
 #[cfg(test)]

src/expr/impl/src/scalar/format.rs

@@ -23,11 +23,25 @@ use thiserror_ext::AsReport;
 use super::string::quote_ident;
 
 /// Formats arguments according to a format string.
+///
+/// # Example
+///
+/// ```slt
+/// query T
+/// select format('%s %s', 'Hello', 'World');
+/// ----
+/// Hello World
+///
+/// query T
+/// select format('%s %s', variadic array['Hello', 'World']);
+/// ----
+/// Hello World
+/// ```
 #[function(
-    "format(varchar, ...) -> varchar",
+    "format(varchar, variadic anyarray) -> varchar",
     prebuild = "Formatter::from_str($0).map_err(|e| ExprError::Parse(e.to_report_string().into()))?"
 )]
-fn format(formatter: &Formatter, row: impl Row, writer: &mut impl Write) -> Result<()> {
+fn format(row: impl Row, formatter: &Formatter, writer: &mut impl Write) -> Result<()> {
     let mut args = row.iter();
     for node in &formatter.nodes {
         match node {

src/expr/impl/src/scalar/jsonb_access.rs

@@ -14,7 +14,8 @@
 
 use std::fmt::Write;
 
-use risingwave_common::types::{JsonbRef, ListRef};
+use risingwave_common::row::Row;
+use risingwave_common::types::JsonbRef;
 use risingwave_expr::function;
 
 /// Extracts JSON object field with the given key.
@@ -91,9 +92,14 @@ pub fn jsonb_array_element(v: JsonbRef<'_>, p: i32) -> Option<JsonbRef<'_>> {
 /// select jsonb_extract_path('{"a": {"b": ["foo","bar"]}}', 'a', 'b', '1');
 /// ----
 /// "bar"
+///
+/// query T
+/// select jsonb_extract_path('{"a": {"b": ["foo","bar"]}}', variadic array['a', 'b', '1']);
+/// ----
+/// "bar"
 /// ```
-#[function("jsonb_extract_path(jsonb, varchar[]) -> jsonb")]
-pub fn jsonb_extract_path<'a>(v: JsonbRef<'a>, path: ListRef<'_>) -> Option<JsonbRef<'a>> {
+#[function("jsonb_extract_path(jsonb, variadic varchar[]) -> jsonb")]
+pub fn jsonb_extract_path(v: JsonbRef<'_>, path: impl Row) -> Option<JsonbRef<'_>> {
     let mut jsonb = v;
     for key in path.iter() {
         // return null if any element is null
@@ -192,11 +198,16 @@ pub fn jsonb_array_element_str(v: JsonbRef<'_>, p: i32, writer: &mut impl Write)
 /// select jsonb_extract_path_text('{"a": {"b": ["foo","bar"]}}', 'a', 'b', '1');
 /// ----
 /// bar
+///
+/// query T
+/// select jsonb_extract_path_text('{"a": {"b": ["foo","bar"]}}', variadic array['a', 'b', '1']);
+/// ----
+/// bar
 /// ```
-#[function("jsonb_extract_path_text(jsonb, varchar[]) -> varchar")]
+#[function("jsonb_extract_path_text(jsonb, variadic varchar[]) -> varchar")]
 pub fn jsonb_extract_path_text(
     v: JsonbRef<'_>,
-    path: ListRef<'_>,
+    path: impl Row,
     writer: &mut impl Write,
 ) -> Option<()> {
     let jsonb = jsonb_extract_path(v, path)?;