Merge pull request #35290 from bigo-sg/function_enumerate_streams

Add function getTypeSerializationStreams

docs/en/sql-reference/functions/other-functions.md

@@ -2499,3 +2499,41 @@ Result:
 │                      286 │
 └──────────────────────────┘
 ```
+
+## getTypeSerializationStreams {#getTypeSerializationStreams}
+
+return the serialization streams of data type.
+
+**Syntax**
+``` sql
+getTypeSerializationStreams(type_name)
+
+getTypeSerializationStreams(column)
+```
+
+**Arguments**
+- `type_name` - Name of data type to get its serialization paths. [String](../../sql-reference/data-types/string.md#string).
+- `column`       - any column which has a data type
+
+**Returned value**
+- List of serialization streams;
+
+Type: [Array](../../sql-reference/data-types/array.md)([String](../../sql-reference/data-types/string.md)).
+
+
+
+**Example**
+
+Query:
+
+``` sql
+SELECT getTypeSerializationStreams('Array(Array(Int8))')
+```
+
+Result:
+
+``` text
+┌───────────────────────getTypeSerializationStreams('Array(Array(Int8))')─────────────────────────────┐
+│ ['{ArraySizes}','{ArrayElements, ArraySizes}','{ArrayElements, ArrayElements, Regular}']            │
+└─────────────────────────────────────────────────────────────────────────────────────────────────────┘
+```

src/Functions/getTypeSerializationStreams.cpp

@@ -0,0 +1,87 @@
+#include <Functions/IFunction.h>
+#include <Functions/FunctionFactory.h>
+#include <DataTypes/DataTypeString.h>
+#include <DataTypes/DataTypeArray.h>
+#include <DataTypes/DataTypeFactory.h>
+#include <DataTypes/Serializations/ISerialization.h>
+#include <Core/Field.h>
+#include <Columns/ColumnString.h>
+#include <Functions/FunctionHelpers.h>
+
+namespace DB
+{
+
+namespace
+{
+
+/// Enumerate stream paths of data type.
+class FunctionGetTypeSerializationStreams : public IFunction
+{
+public:
+    static constexpr auto name = "getTypeSerializationStreams";
+    static FunctionPtr create(ContextPtr)
+    {
+        return std::make_shared<FunctionGetTypeSerializationStreams>();
+    }
+
+    String getName() const override
+    {
+        return name;
+    }
+
+    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
+
+    size_t getNumberOfArguments() const override
+    {
+        return 1;
+    }
+
+    DataTypePtr getReturnTypeImpl(const DataTypes &) const override
+    {
+        return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
+    }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
+    {
+        auto type = getType(arguments[0]);
+
+        SerializationPtr serialization = type->getDefaultSerialization();
+        auto col_res = ColumnArray::create(ColumnString::create());
+        ColumnString & col_res_strings = typeid_cast<ColumnString &>(col_res->getData());
+        ColumnVectorHelper::Offsets & col_res_offsets = typeid_cast<ColumnArray::Offsets &>(col_res->getOffsets());
+        serialization->enumerateStreams([&](const ISerialization::SubstreamPath & substream_path)
+        {
+            col_res_strings.insert(substream_path.toString());
+        });
+        col_res_offsets.push_back(col_res_strings.size());
+        return ColumnConst::create(std::move(col_res), input_rows_count);
+    }
+
+private:
+    static DataTypePtr getType(const ColumnWithTypeAndName & argument)
+    {
+        const IColumn * arg_column = argument.column.get();
+        const ColumnString * arg_string = checkAndGetColumnConstData<ColumnString>(arg_column);
+        if (!arg_string)
+            return argument.type;
+
+        try
+        {
+            DataTypePtr type = DataTypeFactory::instance().get(arg_string->getDataAt(0).toString());
+            return type;
+        }
+        catch (const DB::Exception &)
+        {
+            return argument.type;
+        }
+    }
+};
+
+}
+
+void registerFunctionGetTypeSerializationStreams(FunctionFactory & factory)
+{
+    factory.registerFunction<FunctionGetTypeSerializationStreams>();
+}
+
+}

src/Functions/registerFunctionsMiscellaneous.cpp

@@ -71,15 +71,16 @@ void registerFunctionErrorCodeToName(FunctionFactory &);
 void registerFunctionTcpPort(FunctionFactory &);
 void registerFunctionGetServerPort(FunctionFactory &);
 void registerFunctionByteSize(FunctionFactory &);
-void registerFunctionFile(FunctionFactory & factory);
-void registerFunctionConnectionId(FunctionFactory & factory);
-void registerFunctionPartitionId(FunctionFactory & factory);
+void registerFunctionFile(FunctionFactory &);
+void registerFunctionConnectionId(FunctionFactory &);
+void registerFunctionPartitionId(FunctionFactory &);
 void registerFunctionIsIPAddressContainedIn(FunctionFactory &);
-void registerFunctionQueryID(FunctionFactory & factory);
-void registerFunctionInitialQueryID(FunctionFactory & factory);
+void registerFunctionQueryID(FunctionFactory &);
+void registerFunctionInitialQueryID(FunctionFactory &);
 void registerFunctionServerUUID(FunctionFactory &);
 void registerFunctionZooKeeperSessionUptime(FunctionFactory &);
 void registerFunctionGetOSKernelVersion(FunctionFactory &);
+void registerFunctionGetTypeSerializationStreams(FunctionFactory &);
 void registerFunctionFlattenTuple(FunctionFactory &);
 
 #if USE_ICU
@@ -167,6 +168,7 @@ void registerFunctionsMiscellaneous(FunctionFactory & factory)
     registerFunctionServerUUID(factory);
     registerFunctionZooKeeperSessionUptime(factory);
     registerFunctionGetOSKernelVersion(factory);
+    registerFunctionGetTypeSerializationStreams(factory);
     registerFunctionFlattenTuple(factory);
 
 #if USE_ICU

tests/queries/0_stateless/02240_get_type_serialization_streams.reference

@@ -0,0 +1,8 @@
+['{ArraySizes}','{ArrayElements, Regular}']
+['{ArraySizes}','{ArrayElements, TupleElement(keys, escape_tuple_delimiter = true), Regular}','{ArrayElements, TupleElement(values, escape_tuple_delimiter = true), Regular}']
+['{TupleElement(1, escape_tuple_delimiter = true), Regular}','{TupleElement(2, escape_tuple_delimiter = true), Regular}','{TupleElement(3, escape_tuple_delimiter = true), Regular}']
+['{DictionaryKeys, Regular}','{DictionaryIndexes}']
+['{NullMap}','{NullableElements, Regular}']
+['{ArraySizes}','{ArrayElements, Regular}']
+['{ArraySizes}','{ArrayElements, TupleElement(keys, escape_tuple_delimiter = true), Regular}','{ArrayElements, TupleElement(values, escape_tuple_delimiter = true), Regular}']
+['{TupleElement(1, escape_tuple_delimiter = true), Regular}','{TupleElement(2, escape_tuple_delimiter = true), Regular}','{TupleElement(3, escape_tuple_delimiter = true), Regular}','{TupleElement(4, escape_tuple_delimiter = true), Regular}']

tests/queries/0_stateless/02240_get_type_serialization_streams.sql

@@ -0,0 +1,8 @@
+select getTypeSerializationStreams('Array(Int8)');
+select getTypeSerializationStreams('Map(String, Int64)');
+select getTypeSerializationStreams('Tuple(String, Int64, Float64)');
+select getTypeSerializationStreams('LowCardinality(String)');
+select getTypeSerializationStreams('Nullable(String)');
+select getTypeSerializationStreams([1,2,3]);
+select getTypeSerializationStreams(map('a', 1, 'b', 2));
+select getTypeSerializationStreams(tuple('a', 1, 'b', 2));