fix test

Signed-off-by: Runji Wang <wangrunji0408@163.com>
risingwavelabs · Mar 10, 2023 · 36ef0fe · 36ef0fe
1 parent c37fa65
commit 36ef0fe
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 14 deletions.
diff --git a/src/common/src/types/mod.rs b/src/common/src/types/mod.rs
@@ -347,6 +347,13 @@ impl DataType {
         )
     }
 
+    pub fn as_struct(&self) -> &StructType {
+        match self {
+            DataType::Struct(t) => t,
+            _ => panic!("expect struct type"),
+        }
+    }
+
     /// WARNING: Currently this should only be used in `WatermarkFilterExecutor`. Please be careful
     /// if you want to use this.
     pub fn min(&self) -> ScalarImpl {

diff --git a/src/frontend/src/handler/create_function.rs b/src/frontend/src/handler/create_function.rs
@@ -84,15 +84,22 @@ pub async fn handle_create_function(
             Kind::Scalar(ScalarFunction {})
         }
         Some(CreateFunctionReturns::Table(columns)) => {
-            let datatypes = columns
-                .iter()
-                .map(|c| bind_data_type(&c.data_type))
-                .collect::<Result<Vec<_>>>()?;
-            let names = columns
-                .iter()
-                .map(|c| c.name.real_value())
-                .collect::<Vec<_>>();
-            return_type = DataType::new_struct(datatypes, names);
+            if columns.len() == 1 {
+                // return type is the original type for compatibility
+                // TODO(wrj): should it always be a struct?
+                return_type = bind_data_type(&columns[0].data_type)?;
+            } else {
+                // return type is a struct for multiple columns
+                let datatypes = columns
+                    .iter()
+                    .map(|c| bind_data_type(&c.data_type))
+                    .collect::<Result<Vec<_>>>()?;
+                let names = columns
+                    .iter()
+                    .map(|c| c.name.real_value())
+                    .collect::<Vec<_>>();
+                return_type = DataType::new_struct(datatypes, names);
+            }
             Kind::Table(TableFunction {})
         }
         None => {
@@ -137,11 +144,24 @@ pub async fn handle_create_function(
             .map(|t| arrow_schema::Field::new("", t.into(), true))
             .collect(),
     );
-    let returns = arrow_schema::Schema::new(vec![arrow_schema::Field::new(
-        "",
-        return_type.clone().into(),
-        true,
-    )]);
+    let returns = match kind {
+        Kind::Scalar(_) => arrow_schema::Schema::new(vec![arrow_schema::Field::new(
+            "",
+            return_type.clone().into(),
+            true,
+        )]),
+        Kind::Table(_) => arrow_schema::Schema::new(match &return_type {
+            DataType::Struct(s) => (s.fields.iter())
+                .map(|t| arrow_schema::Field::new("", t.clone().into(), true))
+                .collect(),
+            _ => vec![arrow_schema::Field::new(
+                "",
+                return_type.clone().into(),
+                true,
+            )],
+        }),
+        _ => unreachable!(),
+    };
     client
         .check(&identifier, &args, &returns)
         .await