From d0fc3d9d4ea734b354304b00f7d093eb43378f8e Mon Sep 17 00:00:00 2001
From: fhbai <fhbai@tencent.com>
Date: Wed, 1 Feb 2023 11:17:14 +0800
Subject: [PATCH] allow a three-argument version for table function format

---
 src/TableFunctions/TableFunctionFormat.cpp    | 21 +++++++----
 src/TableFunctions/TableFunctionFormat.h      |  1 +
 .../02542_table_function_format.reference     | 22 ++++++++++++
 .../02542_table_function_format.sql           | 36 +++++++++++++++++++
 4 files changed, 73 insertions(+), 7 deletions(-)
 create mode 100644 tests/queries/0_stateless/02542_table_function_format.reference
 create mode 100644 tests/queries/0_stateless/02542_table_function_format.sql

diff --git a/src/TableFunctions/TableFunctionFormat.cpp b/src/TableFunctions/TableFunctionFormat.cpp
index f2a92b41560..9cc04c2628d 100644
--- a/src/TableFunctions/TableFunctionFormat.cpp
+++ b/src/TableFunctions/TableFunctionFormat.cpp
@@ -4,6 +4,7 @@
 
 #include <Interpreters/Context.h>
 #include <Interpreters/evaluateConstantExpression.h>
+#include <Interpreters/parseColumnsListForTableFunction.h>
 
 #include <Parsers/ASTLiteral.h>
 
@@ -38,23 +39,29 @@ void TableFunctionFormat::parseArguments(const ASTPtr & ast_function, ContextPtr
 
     ASTs & args = args_func.at(0)->children;
 
-    if (args.size() != 2)
-        throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Table function '{}' requires 2 arguments: format and data", getName());
+    if (args.size() != 2 && args.size() != 3)
+        throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH, "Table function '{}' requires 2 or 3 arguments: format [structure] data", getName());
 
     for (auto & arg : args)
         arg = evaluateConstantExpressionOrIdentifierAsLiteral(arg, context);
 
     format = checkAndGetLiteralArgument<String>(args[0], "format");
-    data = checkAndGetLiteralArgument<String>(args[1], "data");
+    data = checkAndGetLiteralArgument<String>(args.back(), "data");
+    if (args.size() == 3)
+        structure = checkAndGetLiteralArgument<String>(args[1], "structure");
 }
 
 ColumnsDescription TableFunctionFormat::getActualTableStructure(ContextPtr context) const
 {
-    ReadBufferIterator read_buffer_iterator = [&](ColumnsDescription &)
+    if (structure == "auto")
     {
-        return std::make_unique<ReadBufferFromString>(data);
-    };
-    return readSchemaFromFormat(format, std::nullopt, read_buffer_iterator, false, context);
+        ReadBufferIterator read_buffer_iterator = [&](ColumnsDescription &)
+        {
+            return std::make_unique<ReadBufferFromString>(data);
+        };
+        return readSchemaFromFormat(format, std::nullopt, read_buffer_iterator, false, context);
+    }
+    return parseColumnsListFromString(structure, context);
 }
 
 Block TableFunctionFormat::parseData(ColumnsDescription columns, ContextPtr context) const
diff --git a/src/TableFunctions/TableFunctionFormat.h b/src/TableFunctions/TableFunctionFormat.h
index c6db322343b..d64ab14cb64 100644
--- a/src/TableFunctions/TableFunctionFormat.h
+++ b/src/TableFunctions/TableFunctionFormat.h
@@ -28,6 +28,7 @@ private:
 
     String format;
     String data;
+    String structure = "auto";
 };
 
 }
diff --git a/tests/queries/0_stateless/02542_table_function_format.reference b/tests/queries/0_stateless/02542_table_function_format.reference
new file mode 100644
index 00000000000..c8488967144
--- /dev/null
+++ b/tests/queries/0_stateless/02542_table_function_format.reference
@@ -0,0 +1,22 @@
+a	Nullable(String)					
+b	Nullable(Int64)					
+a	String					
+b	Int64					
+Hello	111
+World	123
+Hello	111
+World	123
+c1	Nullable(Int64)					
+c2	Nullable(Int64)					
+c3	Array(Nullable(Int64))					
+c4	Array(Array(Nullable(String)))					
+a1	Int32					
+a2	UInt64					
+a3	Array(Int32)					
+a4	Array(Array(String))					
+1	2	[1,2,3]	[['abc'],[],['d','e']]
+1	2	[1,2,3]	[['abc'],[],['d','e']]
+20210129005809043707
+123456789
+987654321
+cust_id	UInt128					
diff --git a/tests/queries/0_stateless/02542_table_function_format.sql b/tests/queries/0_stateless/02542_table_function_format.sql
new file mode 100644
index 00000000000..e32e9001b9f
--- /dev/null
+++ b/tests/queries/0_stateless/02542_table_function_format.sql
@@ -0,0 +1,36 @@
+desc format(JSONEachRow, 
+$$
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+$$);
+
+desc format(JSONEachRow, 'a String, b Int64',
+$$
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+$$);
+
+select * from format(JSONEachRow, 'a String, b Int64',
+$$
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+{"a": "Hello", "b": 111}
+{"a": "World", "b": 123}
+$$);
+
+desc format(CSV, '1,2,"[1,2,3]","[[\'abc\'], [], [\'d\', \'e\']]"');
+desc format(CSV, 'a1 Int32, a2 UInt64, a3 Array(Int32), a4 Array(Array(String))', '1,2,"[1,2,3]","[[\'abc\'], [], [\'d\', \'e\']]"');
+select * from format(CSV, '1,2,"[1,2,3]","[[\'abc\'], [], [\'d\', \'e\']]"');
+select * from format(CSV, 'a1 Int32, a2 UInt64, a3 Array(Int32), a4 Array(Array(String))', '1,2,"[1,2,3]","[[\'abc\'], [], [\'d\', \'e\']]"');
+
+drop table if exists test;
+
+create table test as format(TSV, 'cust_id UInt128', '20210129005809043707\n123456789\n987654321');
+
+select * from test;
+desc table test;
+drop table test;