FLASH-821 support JSON encode in coprocessor

dbms/src/Flash/Coprocessor/ArrowColCodec.cpp

@@ -1,5 +1,3 @@
-#include <Flash/Coprocessor/ArrowColCodec.h>
-
 #include <Columns/ColumnDecimal.h>
 #include <Columns/ColumnNullable.h>
 #include <Columns/ColumnString.h>
@@ -11,6 +9,7 @@
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeString.h>
 #include <DataTypes/DataTypesNumber.h>
+#include <Flash/Coprocessor/ArrowColCodec.h>
 #include <Flash/Coprocessor/DAGUtils.h>
 #include <Functions/FunctionHelpers.h>
 #include <IO/Endian.h>
@@ -336,6 +335,7 @@ void flashColToArrowCol(TiDBColumn & dag_column, const ColumnWithTypeAndName & f
         case TiDB::TypeLongBlob:
         case TiDB::TypeMediumBlob:
         case TiDB::TypeTinyBlob:
+        case TiDB::TypeJSON:
             if (!checkDataType<DataTypeString>(type))
                 throw Exception(
                     "Type un-matched during arrow encode, target col type is string and source column type is " + type->getName(),
@@ -631,6 +631,7 @@ const char * arrowColToFlashCol(const char * pos, UInt8 field_length, UInt32 nul
         case TiDB::TypeTinyBlob:
         case TiDB::TypeMediumBlob:
         case TiDB::TypeLongBlob:
+        case TiDB::TypeJSON:
             return arrowStringColToFlashCol(pos, field_length, null_count, null_bitmap, offsets, flash_col, col_info, length);
         case TiDB::TypeBit:
             return arrowBitColToFlashCol(pos, field_length, null_count, null_bitmap, offsets, flash_col, col_info, length);

dbms/src/Flash/Coprocessor/DAGUtils.cpp

@@ -1,7 +1,6 @@
-#include <Flash/Coprocessor/DAGUtils.h>
-
 #include <Core/Types.h>
 #include <Flash/Coprocessor/DAGCodec.h>
+#include <Flash/Coprocessor/DAGUtils.h>
 #include <Functions/FunctionHelpers.h>
 #include <Interpreters/Context.h>
 #include <Storages/Transaction/Datum.h>
@@ -275,13 +274,17 @@ bool exprHasValidFieldType(const tipb::Expr & expr)
 
 bool isUnsupportedEncodeType(const std::vector<tipb::FieldType> & types, tipb::EncodeType encode_type)
 {
+    const static std::unordered_map<tipb::EncodeType, std::unordered_set<Int32>> unsupported_types_map({
+        {tipb::EncodeType::TypeCHBlock, {TiDB::TypeSet, TiDB::TypeGeometry, TiDB::TypeNull, TiDB::TypeEnum, TiDB::TypeJSON}},
+        {tipb::EncodeType::TypeChunk, {TiDB::TypeSet, TiDB::TypeGeometry, TiDB::TypeNull}},
+    });
+
     if (encode_type == tipb::EncodeType::TypeDefault)
         return false;
     for (const auto & type : types)
     {
-        if (encode_type == tipb::EncodeType::TypeCHBlock && (type.tp() == TiDB::TypeSet || type.tp() == TiDB::TypeEnum))
-            return true;
-        if (encode_type == tipb::EncodeType::TypeChunk && type.tp() == TiDB::TypeSet)
+        auto unsupported_set = unsupported_types_map.find(encode_type);
+        if (unsupported_set != unsupported_types_map.end() && unsupported_set->second.find(type.tp()) != unsupported_set->second.end())
             return true;
     }
     return false;
@@ -319,6 +322,7 @@ UInt8 getFieldLengthForArrowEncode(Int32 tp)
         case TiDB::TypeLongBlob:
         case TiDB::TypeBit:
         case TiDB::TypeEnum:
+        case TiDB::TypeJSON:
             return VAR_SIZE;
         default:
             throw Exception("not supported field type in arrow encode: " + std::to_string(tp));

dbms/src/Storages/Transaction/Codec.cpp

@@ -381,6 +381,12 @@ void EncodeCompactBytes(const String & str, std::stringstream & ss)
     ss.write(str.c_str(), str.size());
 }
 
+void EncodeJSON(const String & str, std::stringstream & ss)
+{
+    // TiFlash store the JSON binary as string, so just return the string
+    ss.write(str.c_str(), str.size());
+}
+
 void EncodeVarUInt(UInt64 num, std::stringstream & ss) { TiKV::writeVarUInt(num, ss); }
 
 void EncodeVarInt(Int64 num, std::stringstream & ss) { TiKV::writeVarInt(num, ss); }
@@ -517,10 +523,7 @@ void EncodeDecimal(const Field & field, std::stringstream & ss)
 void EncodeDatum(const Field & field, TiDB::CodecFlag flag, std::stringstream & ss)
 {
     if (field.isNull())
-    {
-        ss << UInt8(TiDB::CodecFlagNil);
-        return;
-    }
+        flag = TiDB::CodecFlagNil;
     ss << UInt8(flag);
     switch (flag)
     {
@@ -540,6 +543,10 @@ void EncodeDatum(const Field & field, TiDB::CodecFlag flag, std::stringstream &
             return EncodeVarUInt(field.safeGet<UInt64>(), ss);
         case TiDB::CodecFlagDuration:
             return EncodeInt64(field.safeGet<Int64>(), ss);
+        case TiDB::CodecFlagJson:
+            return EncodeJSON(field.safeGet<String>(), ss);
+        case TiDB::CodecFlagNil:
+            return;
         default:
             throw Exception("Not implemented codec flag: " + std::to_string(flag), ErrorCodes::LOGICAL_ERROR);
     }

dbms/src/Storages/Transaction/Codec.h

@@ -58,6 +58,8 @@ void EncodeBytes(const String & ori_str, std::stringstream & ss);
 
 void EncodeCompactBytes(const String & str, std::stringstream & ss);
 
+void EncodeJSON(const String & str, std::stringstream & ss);
+
 void EncodeVarUInt(UInt64 num, std::stringstream & ss);
 
 void EncodeVarInt(Int64 num, std::stringstream & ss);