support multi parition columns

lgbo-ustc · lgbo-ustc · commit 0da3ff4aa0ef · 2022-09-16T09:34:14.000+08:00
diff --git a/utils/local-engine/Operator/PartitionColumnFillingTransform.cpp b/utils/local-engine/Operator/PartitionColumnFillingTransform.cpp
@@ -6,12 +6,16 @@
 #include <IO/ReadBufferFromString.h>
 #include <IO/ReadHelpers.h>
 #include <Common/StringUtils.h>
+#include "Processors/Chunk.h"
 #include <Columns/IColumn.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/Serializations/ISerialization.h>
 #include <Interpreters/DatabaseAndTableWithAlias.h>
 #include <base/DayNum.h>
 
+#include <Poco/Logger.h>
+#include <base/logger_useful.h>
+
 
 using namespace DB;
 
@@ -26,40 +30,37 @@ namespace ErrorCodes
 namespace local_engine
 {
 template <typename Type>
-requires(
-    std::is_same_v<Type, Int8> || std::is_same_v<Type, UInt16> || std::is_same_v<Type, Int16> || std::is_same_v<Type, Int32> || std::is_same_v<Type, Int64>)
-    ColumnPtr createIntPartitionColumn(DataTypePtr column_type, std::string partition_value)
+    requires(std::is_same_v<Type, Int8> || std::is_same_v<Type, UInt16> || std::is_same_v<Type, Int16> || std::is_same_v<Type, Int32> || std::is_same_v<Type, Int64>)
+ColumnPtr createIntPartitionColumn(DataTypePtr column_type, std::string partition_value, size_t rows)
 {
     Type value;
     auto value_buffer = ReadBufferFromString(partition_value);
     readIntText(value, value_buffer);
-    return column_type->createColumnConst(1, value);
+    return column_type->createColumnConst(rows, value);
 }
 
 template <typename Type>
-requires(std::is_same_v<Type, Float32> || std::is_same_v<Type, Float64>) ColumnPtr
-    createFloatPartitionColumn(DataTypePtr column_type, std::string partition_value)
+    requires(std::is_same_v<Type, Float32> || std::is_same_v<Type, Float64>)
+ColumnPtr createFloatPartitionColumn(DataTypePtr column_type, std::string partition_value, size_t rows)
 {
     Type value;
     auto value_buffer = ReadBufferFromString(partition_value);
     readFloatText(value, value_buffer);
-    return column_type->createColumnConst(1, value);
+    return column_type->createColumnConst(rows, value);
 }
 
-//template <>
-//ColumnPtr createFloatPartitionColumn<Float32>(DataTypePtr column_type, std::string partition_value);
-//template <>
-//ColumnPtr createFloatPartitionColumn<Float64>(DataTypePtr column_type, std::string partition_value);
 
 PartitionColumnFillingTransform::PartitionColumnFillingTransform(
-    const DB::Block & input_, const DB::Block & output_, const String & partition_col_name_, const String & partition_col_value_)
-    : ISimpleTransform(input_, output_, true), partition_col_name(partition_col_name_), partition_col_value(partition_col_value_)
+    const DB::Block & input_, const DB::Block & output_, const PartitionValues & partition_columns_)
+    : ISimpleTransform(input_, output_, true), partition_column_values(partition_columns_)
 {
-    partition_col_type = output_.getByName(partition_col_name_).type;
-    partition_column = createPartitionColumn();
+    for (const auto & value : partition_column_values)
+    {
+        partition_columns[value.first] = value.second;
+    }
 }
 
-/// In the case that a partition column is wrapper by nullable and LowCardinality, we need to keep the data type same
+/// In the case that a partition column is wrapper by nullable or LowCardinality, we need to keep the data type same
 /// as input.
 ColumnPtr PartitionColumnFillingTransform::tryWrapPartitionColumn(const ColumnPtr & nested_col, DataTypePtr original_data_type)
 {
@@ -71,9 +72,10 @@ ColumnPtr PartitionColumnFillingTransform::tryWrapPartitionColumn(const ColumnPt
     return result;
 }
 
-ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()
+ColumnPtr PartitionColumnFillingTransform::createPartitionColumn(const String & parition_col, const String & partition_col_value, size_t rows)
 {
     ColumnPtr result;
+    auto partition_col_type = output.getHeader().getByName(parition_col).type;
     DataTypePtr nested_type = partition_col_type;
     if (const DataTypeNullable * nullable_type = checkAndGetDataType<DataTypeNullable>(partition_col_type.get()))
     {
@@ -86,45 +88,45 @@ ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()
     WhichDataType which(nested_type);
     if (which.isInt8())
     {
-        result = createIntPartitionColumn<Int8>(nested_type, partition_col_value);
+        result = createIntPartitionColumn<Int8>(nested_type, partition_col_value, rows);
     }
     else if (which.isInt16())
     {
-        result = createIntPartitionColumn<Int16>(nested_type, partition_col_value);
+        result = createIntPartitionColumn<Int16>(nested_type, partition_col_value, rows);
     }
     else if (which.isInt32())
     {
-        result = createIntPartitionColumn<Int32>(nested_type, partition_col_value);
+        result = createIntPartitionColumn<Int32>(nested_type, partition_col_value, rows);
     }
     else if (which.isInt64())
     {
-        result = createIntPartitionColumn<Int64>(nested_type, partition_col_value);
+        result = createIntPartitionColumn<Int64>(nested_type, partition_col_value, rows);
     }
     else if (which.isFloat32())
     {
-        result = createFloatPartitionColumn<Float32>(nested_type, partition_col_value);
+        result = createFloatPartitionColumn<Float32>(nested_type, partition_col_value, rows);
     }
     else if (which.isFloat64())
     {
-        result = createFloatPartitionColumn<Float64>(nested_type, partition_col_value);
+        result = createFloatPartitionColumn<Float64>(nested_type, partition_col_value, rows);
     }
     else if (which.isDate())
     {
         DayNum value;
         auto value_buffer = ReadBufferFromString(partition_col_value);
         readDateText(value, value_buffer);
-        result = nested_type->createColumnConst(1, value);
+        result = nested_type->createColumnConst(rows, value);
     }
     else if (which.isDate32())
     {
         ExtendedDayNum value;
         auto value_buffer = ReadBufferFromString(partition_col_value);
         readDateText(value, value_buffer);
-        result = nested_type->createColumnConst(1, value.toUnderType());
+        result = nested_type->createColumnConst(rows, value.toUnderType());
     }
     else if (which.isString())
     {
-        result = nested_type->createColumnConst(1, partition_col_value);
+        result = nested_type->createColumnConst(rows, partition_col_value);
     }
     else
     {
@@ -136,14 +138,30 @@ ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()
 
 void PartitionColumnFillingTransform::transform(DB::Chunk & chunk)
 {
-    size_t partition_column_position = output.getHeader().getPositionByName(partition_col_name);
-    if (partition_column_position == input.getHeader().columns())
-    {
-        chunk.addColumn(partition_column->cloneResized(chunk.getNumRows()));
-    }
-    else
+    auto rows = chunk.getNumRows();
+    auto input_cols = chunk.detachColumns();
+    Columns result_cols;
+    auto input_header = input.getHeader();
+    for (const auto & output_col : output.getHeader())
     {
-        chunk.addColumn(partition_column_position, partition_column->cloneResized(chunk.getNumRows()));
+        if (input_header.has(output_col.name))
+        {
+            size_t pos = input_header.getPositionByName(output_col.name);
+            result_cols.push_back(input_cols[pos]);
+        }
+        else
+        {
+            // it's a partition column
+            auto it = partition_columns.find(output_col.name);
+            if (it == partition_columns.end())
+            {
+                throw Exception(ErrorCodes::LOGICAL_ERROR, "Not found column({}) in parition columns", output_col.name);
+            }
+            result_cols.emplace_back(createPartitionColumn(it->first, it->second, rows));
+
+        }
+        
     }
+    chunk = DB::Chunk(std::move(result_cols), rows);
 }
 }
diff --git a/utils/local-engine/Operator/PartitionColumnFillingTransform.h b/utils/local-engine/Operator/PartitionColumnFillingTransform.h
@@ -1,6 +1,10 @@
 #pragma once
 
 #include <Processors/ISimpleTransform.h>
+#include "Common/StringUtils.h"
+#include "Columns/IColumn.h"
+#include "Core/Block.h"
+#include "DataTypes/Serializations/ISerialization.h"
 
 namespace local_engine
 {
@@ -10,22 +14,19 @@ class PartitionColumnFillingTransform : public DB::ISimpleTransform
     PartitionColumnFillingTransform(
         const DB::Block & input_,
         const DB::Block & output_,
-        const String & partition_col_name_,
-        const String & partition_col_value_);
+        const PartitionValues & partition_columns_);
     void transform(DB::Chunk & chunk) override;
     String getName() const override
     {
         return "PartitionColumnFillingTransform";
     }
 
 private:
-    DB::ColumnPtr createPartitionColumn();
+    DB::ColumnPtr createPartitionColumn(const String & parition_col, const String & partition_col_value, size_t row);
     static DB::ColumnPtr tryWrapPartitionColumn(const DB::ColumnPtr & nested_col, DB::DataTypePtr original_data_type);
 
-    DB::DataTypePtr partition_col_type;
-    String partition_col_name;
-    String partition_col_value;
-    DB::ColumnPtr partition_column;
+    PartitionValues partition_column_values;
+    std::map<String, String> partition_columns;
 };
 
 }
diff --git a/utils/local-engine/Parser/SerializedPlanParser.cpp b/utils/local-engine/Parser/SerializedPlanParser.cpp
@@ -1,4 +1,5 @@
-#include <base/logger_useful.h>
+#include "SerializedPlanParser.h"
+#include <memory>
 #include <AggregateFunctions/AggregateFunctionFactory.h>
 #include <AggregateFunctions/registerAggregateFunctions.h>
 #include <Builder/BroadCastJoinBuilder.h>
@@ -41,7 +42,7 @@
 #include <Common/MergeTreeTool.h>
 #include <Common/StringUtils.h>
 
-#include "SerializedPlanParser.h"
+#include <google/protobuf/util/json_util.h>
 
 namespace DB
 {
@@ -197,19 +198,14 @@ QueryPlanPtr SerializedPlanParser::parseReadRealWithLocalFile(const substrait::R
     }
     auto header = parseNameStruct(rel.base_schema());
     PartitionValues partition_values = StringUtils::parsePartitionTablePath(files_info->files[0]);
-    if (partition_values.size() > 1)
-    {
-        throw Exception(ErrorCodes::NOT_IMPLEMENTED, "doesn't support multiple level partition.");
-    }
-    ProcessorPtr partition_transform;
-    if (!partition_values.empty())
+    
+    auto origin_header = header.cloneEmpty();
+    for (const auto & partition_value : partition_values)
     {
-        auto origin_header = header.cloneEmpty();
-        PartitionValue partition_value = partition_values[0];
         header.erase(partition_value.first);
-        partition_transform
-            = std::make_shared<PartitionColumnFillingTransform>(header, origin_header, partition_value.first, partition_value.second);
     }
+    ProcessorPtr partition_transform = std::make_shared<PartitionColumnFillingTransform>(header, origin_header, partition_values);
+
     auto query_plan = std::make_unique<QueryPlan>();
     std::shared_ptr<IProcessor> source = std::make_shared<BatchParquetFileSource>(files_info, header, context);
     auto source_pipe = Pipe(source);
@@ -1281,7 +1277,18 @@ QueryPlanPtr SerializedPlanParser::parse(std::string & plan)
 {
     auto plan_ptr = std::make_unique<substrait::Plan>();
     plan_ptr->ParseFromString(plan);
-    LOG_DEBUG(&Poco::Logger::get("SerializedPlanParser"), "parse plan \n{}", plan_ptr->DebugString());
+
+    auto printPlan = [](const std::string & plan_raw){
+        substrait::Plan plan;
+        plan.ParseFromString(plan_raw);
+        std::string json_ret;
+        google::protobuf::util::JsonPrintOptions json_opt;
+        json_opt.add_whitespace = true;
+        google::protobuf::util::MessageToJsonString(plan, &json_ret, json_opt);
+        return json_ret;
+    };
+
+    LOG_DEBUG(&Poco::Logger::get("SerializedPlanParser"), "parse plan \n{}", printPlan(plan));
     return parse(std::move(plan_ptr));
 }
 void SerializedPlanParser::initFunctionEnv()

Original file line number	Diff line number	Diff line change
`@@ -6,12 +6,16 @@`
`6`	`6`	`#include <IO/ReadBufferFromString.h>`
`7`	`7`	`#include <IO/ReadHelpers.h>`
`8`	`8`	`#include <Common/StringUtils.h>`
	`9`	`+#include "Processors/Chunk.h"`
`9`	`10`	`#include <Columns/IColumn.h>`
`10`	`11`	`#include <DataTypes/DataTypeNullable.h>`
`11`	`12`	`#include <DataTypes/Serializations/ISerialization.h>`
`12`	`13`	`#include <Interpreters/DatabaseAndTableWithAlias.h>`
`13`	`14`	`#include <base/DayNum.h>`
`14`	`15`
	`16`	`+#include <Poco/Logger.h>`
	`17`	`+#include <base/logger_useful.h>`
	`18`	`+`
`15`	`19`
`16`	`20`	`using namespace DB;`
`17`	`21`
`@@ -26,40 +30,37 @@ namespace ErrorCodes`
`26`	`30`	`namespace local_engine`
`27`	`31`	`{`
`28`	`32`	`template <typename Type>`
`29`		`-requires(`
`30`		`- std::is_same_v<Type, Int8> \|\| std::is_same_v<Type, UInt16> \|\| std::is_same_v<Type, Int16> \|\| std::is_same_v<Type, Int32> \|\| std::is_same_v<Type, Int64>)`
`31`		`- ColumnPtr createIntPartitionColumn(DataTypePtr column_type, std::string partition_value)`
	`33`	`+ requires(std::is_same_v<Type, Int8> \|\| std::is_same_v<Type, UInt16> \|\| std::is_same_v<Type, Int16> \|\| std::is_same_v<Type, Int32> \|\| std::is_same_v<Type, Int64>)`
	`34`	`+ColumnPtr createIntPartitionColumn(DataTypePtr column_type, std::string partition_value, size_t rows)`
`32`	`35`	`{`
`33`	`36`	`Type value;`
`34`	`37`	`auto value_buffer = ReadBufferFromString(partition_value);`
`35`	`38`	`readIntText(value, value_buffer);`
`36`		`- return column_type->createColumnConst(1, value);`
	`39`	`+ return column_type->createColumnConst(rows, value);`
`37`	`40`	`}`
`38`	`41`
`39`	`42`	`template <typename Type>`
`40`		`-requires(std::is_same_v<Type, Float32> \|\| std::is_same_v<Type, Float64>) ColumnPtr`
`41`		`- createFloatPartitionColumn(DataTypePtr column_type, std::string partition_value)`
	`43`	`+ requires(std::is_same_v<Type, Float32> \|\| std::is_same_v<Type, Float64>)`
	`44`	`+ColumnPtr createFloatPartitionColumn(DataTypePtr column_type, std::string partition_value, size_t rows)`
`42`	`45`	`{`
`43`	`46`	`Type value;`
`44`	`47`	`auto value_buffer = ReadBufferFromString(partition_value);`
`45`	`48`	`readFloatText(value, value_buffer);`
`46`		`- return column_type->createColumnConst(1, value);`
	`49`	`+ return column_type->createColumnConst(rows, value);`
`47`	`50`	`}`
`48`	`51`
`49`		`-//template <>`
`50`		`-//ColumnPtr createFloatPartitionColumn<Float32>(DataTypePtr column_type, std::string partition_value);`
`51`		`-//template <>`
`52`		`-//ColumnPtr createFloatPartitionColumn<Float64>(DataTypePtr column_type, std::string partition_value);`
`53`	`52`
`54`	`53`	`PartitionColumnFillingTransform::PartitionColumnFillingTransform(`
`55`		`- const DB::Block & input_, const DB::Block & output_, const String & partition_col_name_, const String & partition_col_value_)`
`56`		`- : ISimpleTransform(input_, output_, true), partition_col_name(partition_col_name_), partition_col_value(partition_col_value_)`
	`54`	`+ const DB::Block & input_, const DB::Block & output_, const PartitionValues & partition_columns_)`
	`55`	`+ : ISimpleTransform(input_, output_, true), partition_column_values(partition_columns_)`
`57`	`56`	`{`
`58`		`- partition_col_type = output_.getByName(partition_col_name_).type;`
`59`		`- partition_column = createPartitionColumn();`
	`57`	`+ for (const auto & value : partition_column_values)`
	`58`	`+ {`
	`59`	`+ partition_columns[value.first] = value.second;`
	`60`	`+ }`
`60`	`61`	`}`
`61`	`62`
`62`		`-/// In the case that a partition column is wrapper by nullable and LowCardinality, we need to keep the data type same`
	`63`	`+/// In the case that a partition column is wrapper by nullable or LowCardinality, we need to keep the data type same`
`63`	`64`	`/// as input.`
`64`	`65`	`ColumnPtr PartitionColumnFillingTransform::tryWrapPartitionColumn(const ColumnPtr & nested_col, DataTypePtr original_data_type)`
`65`	`66`	`{`
`@@ -71,9 +72,10 @@ ColumnPtr PartitionColumnFillingTransform::tryWrapPartitionColumn(const ColumnPt`
`71`	`72`	`return result;`
`72`	`73`	`}`
`73`	`74`
`74`		`-ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()`
	`75`	`+ColumnPtr PartitionColumnFillingTransform::createPartitionColumn(const String & parition_col, const String & partition_col_value, size_t rows)`
`75`	`76`	`{`
`76`	`77`	`ColumnPtr result;`
	`78`	`+ auto partition_col_type = output.getHeader().getByName(parition_col).type;`
`77`	`79`	`DataTypePtr nested_type = partition_col_type;`
`78`	`80`	`if (const DataTypeNullable * nullable_type = checkAndGetDataType<DataTypeNullable>(partition_col_type.get()))`
`79`	`81`	`{`
`@@ -86,45 +88,45 @@ ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()`
`86`	`88`	`WhichDataType which(nested_type);`
`87`	`89`	`if (which.isInt8())`
`88`	`90`	`{`
`89`		`- result = createIntPartitionColumn<Int8>(nested_type, partition_col_value);`
	`91`	`+ result = createIntPartitionColumn<Int8>(nested_type, partition_col_value, rows);`
`90`	`92`	`}`
`91`	`93`	`else if (which.isInt16())`
`92`	`94`	`{`
`93`		`- result = createIntPartitionColumn<Int16>(nested_type, partition_col_value);`
	`95`	`+ result = createIntPartitionColumn<Int16>(nested_type, partition_col_value, rows);`
`94`	`96`	`}`
`95`	`97`	`else if (which.isInt32())`
`96`	`98`	`{`
`97`		`- result = createIntPartitionColumn<Int32>(nested_type, partition_col_value);`
	`99`	`+ result = createIntPartitionColumn<Int32>(nested_type, partition_col_value, rows);`
`98`	`100`	`}`
`99`	`101`	`else if (which.isInt64())`
`100`	`102`	`{`
`101`		`- result = createIntPartitionColumn<Int64>(nested_type, partition_col_value);`
	`103`	`+ result = createIntPartitionColumn<Int64>(nested_type, partition_col_value, rows);`
`102`	`104`	`}`
`103`	`105`	`else if (which.isFloat32())`
`104`	`106`	`{`
`105`		`- result = createFloatPartitionColumn<Float32>(nested_type, partition_col_value);`
	`107`	`+ result = createFloatPartitionColumn<Float32>(nested_type, partition_col_value, rows);`
`106`	`108`	`}`
`107`	`109`	`else if (which.isFloat64())`
`108`	`110`	`{`
`109`		`- result = createFloatPartitionColumn<Float64>(nested_type, partition_col_value);`
	`111`	`+ result = createFloatPartitionColumn<Float64>(nested_type, partition_col_value, rows);`
`110`	`112`	`}`
`111`	`113`	`else if (which.isDate())`
`112`	`114`	`{`
`113`	`115`	`DayNum value;`
`114`	`116`	`auto value_buffer = ReadBufferFromString(partition_col_value);`
`115`	`117`	`readDateText(value, value_buffer);`
`116`		`- result = nested_type->createColumnConst(1, value);`
	`118`	`+ result = nested_type->createColumnConst(rows, value);`
`117`	`119`	`}`
`118`	`120`	`else if (which.isDate32())`
`119`	`121`	`{`
`120`	`122`	`ExtendedDayNum value;`
`121`	`123`	`auto value_buffer = ReadBufferFromString(partition_col_value);`
`122`	`124`	`readDateText(value, value_buffer);`
`123`		`- result = nested_type->createColumnConst(1, value.toUnderType());`
	`125`	`+ result = nested_type->createColumnConst(rows, value.toUnderType());`
`124`	`126`	`}`
`125`	`127`	`else if (which.isString())`
`126`	`128`	`{`
`127`		`- result = nested_type->createColumnConst(1, partition_col_value);`
	`129`	`+ result = nested_type->createColumnConst(rows, partition_col_value);`
`128`	`130`	`}`
`129`	`131`	`else`
`130`	`132`	`{`
`@@ -136,14 +138,30 @@ ColumnPtr PartitionColumnFillingTransform::createPartitionColumn()`
`136`	`138`
`137`	`139`	`void PartitionColumnFillingTransform::transform(DB::Chunk & chunk)`
`138`	`140`	`{`
`139`		`- size_t partition_column_position = output.getHeader().getPositionByName(partition_col_name);`
`140`		`- if (partition_column_position == input.getHeader().columns())`
`141`		`- {`
`142`		`- chunk.addColumn(partition_column->cloneResized(chunk.getNumRows()));`
`143`		`- }`
`144`		`- else`
	`141`	`+ auto rows = chunk.getNumRows();`
	`142`	`+ auto input_cols = chunk.detachColumns();`
	`143`	`+ Columns result_cols;`
	`144`	`+ auto input_header = input.getHeader();`
	`145`	`+ for (const auto & output_col : output.getHeader())`
`145`	`146`	`{`
`146`		`- chunk.addColumn(partition_column_position, partition_column->cloneResized(chunk.getNumRows()));`
	`147`	`+ if (input_header.has(output_col.name))`
	`148`	`+ {`
	`149`	`+ size_t pos = input_header.getPositionByName(output_col.name);`
	`150`	`+ result_cols.push_back(input_cols[pos]);`
	`151`	`+ }`
	`152`	`+ else`
	`153`	`+ {`
	`154`	`+ // it's a partition column`
	`155`	`+ auto it = partition_columns.find(output_col.name);`
	`156`	`+ if (it == partition_columns.end())`
	`157`	`+ {`
	`158`	`+ throw Exception(ErrorCodes::LOGICAL_ERROR, "Not found column({}) in parition columns", output_col.name);`
	`159`	`+ }`
	`160`	`+ result_cols.emplace_back(createPartitionColumn(it->first, it->second, rows));`
	`161`	`+`
	`162`	`+ }`
	`163`	`+`
`147`	`164`	`}`
	`165`	`+ chunk = DB::Chunk(std::move(result_cols), rows);`
`148`	`166`	`}`
`149`	`167`	`}`