coprocessor: Refactor encoding TiDBColumn

dbms/src/Columns/ColumnArray.cpp

@@ -1110,25 +1110,9 @@ void ColumnArray::insertFromDatumData(const char * data, size_t length)
     insertData(data, precise_data_size);
 }
 
-size_t ColumnArray::encodeIntoDatumData(size_t element_idx, WriteBuffer & writer) const
+std::pair<UInt32, StringRef> ColumnArray::getElementRef(size_t element_idx) const
 {
-    RUNTIME_CHECK(boost::endian::order::native == boost::endian::order::little);
-
-    RUNTIME_CHECK(checkAndGetColumn<ColumnVector<Float32>>(&getData()));
-    RUNTIME_CHECK(getData().isFixedAndContiguous());
-
-    auto n = static_cast<UInt32>(sizeAt(element_idx));
-
-    writeIntBinary(n, writer);
-    size_t encoded_size = sizeof(UInt32);
-
-    auto data = getDataAt(element_idx);
-    RUNTIME_CHECK(data.size == n * sizeof(Float32));
-    writer.write(data.data, data.size);
-    encoded_size += data.size;
-
-    return encoded_size;
+    return {static_cast<UInt32>(sizeAt(element_idx)), getDataAt(element_idx)};
 }
 
-
 } // namespace DB

dbms/src/Columns/ColumnArray.h

@@ -174,7 +174,7 @@ class ColumnArray final : public COWPtrHelper<IColumn, ColumnArray>
 
     void insertFromDatumData(const char * data, size_t length) override;
 
-    size_t encodeIntoDatumData(size_t element_idx, WriteBuffer & writer) const;
+    std::pair<UInt32, StringRef> getElementRef(size_t element_idx) const;
 
 private:
     ColumnPtr data;

dbms/src/Flash/Coprocessor/ArrowColCodec.cpp

@@ -304,6 +304,10 @@ void flashArrayFloat32ColToArrowCol(
     // We only unwrap the NULLABLE() part.
     const IColumn * nested_col = getNestedCol(flash_col_untyped);
     const auto * flash_col = checkAndGetColumn<ColumnArray>(nested_col);
+
+    RUNTIME_CHECK(checkAndGetColumn<ColumnVector<Float32>>(&flash_col->getData()));
+    RUNTIME_CHECK(flash_col->getData().isFixedAndContiguous());
+
     for (size_t i = start_index; i < end_index; i++)
     {
         // todo check if we can convert flash_col to DAG col directly since the internal representation is almost the same
@@ -316,9 +320,8 @@ void flashArrayFloat32ColToArrowCol(
             }
         }
 
-        auto encoded_size = flash_col->encodeIntoDatumData(i, *dag_column.data);
-        RUNTIME_CHECK(encoded_size > 0);
-        dag_column.finishAppendVar(encoded_size);
+        auto [num_elems, elem_bytes] = flash_col->getElementRef(i);
+        dag_column.appendVectorF32(num_elems, elem_bytes);
     }
 }
 

dbms/src/Flash/Coprocessor/TiDBChunk.cpp

@@ -31,7 +31,7 @@ namespace DB
 {
 TiDBChunk::TiDBChunk(const std::vector<tipb::FieldType> & field_types)
 {
-    for (auto & type : field_types)
+    for (const auto & type : field_types)
     {
         columns.emplace_back(getFieldLengthForArrowEncode(type.tp()));
     }

dbms/src/Flash/Coprocessor/TiDBColumn.cpp

@@ -12,10 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <Common/Exception.h>
 #include <Flash/Coprocessor/DAGCodec.h>
 #include <Flash/Coprocessor/TiDBColumn.h>
 #include <IO/Operators.h>
+#include <IO/WriteHelpers.h>
 #include <IO/copyData.h>
+#include <common/types.h>
 
 namespace DB
 {
@@ -118,6 +121,19 @@ void TiDBColumn::append(const TiDBEnum & ti_enum)
     finishAppendVar(size);
 }
 
+void TiDBColumn::appendVectorF32(UInt32 num_elem, StringRef elem_bytes)
+{
+    encodeLittleEndian<UInt32>(num_elem, *data);
+    size_t encoded_size = sizeof(UInt32);
+
+    RUNTIME_CHECK(elem_bytes.size == num_elem * sizeof(Float32));
+    data->write(elem_bytes.data, elem_bytes.size);
+    encoded_size += elem_bytes.size;
+
+    RUNTIME_CHECK(encoded_size > 0);
+    finishAppendVar(encoded_size);
+}
+
 void TiDBColumn::append(const TiDBBit & bit)
 {
     data->write(bit.val.data, bit.val.size);

dbms/src/Flash/Coprocessor/TiDBColumn.h

@@ -41,20 +41,20 @@ class TiDBColumn
     void append(const TiDBDecimal & decimal);
     void append(const TiDBBit & bit);
     void append(const TiDBEnum & ti_enum);
+    void appendVectorF32(UInt32 num_elem, StringRef elem_bytes);
     void encodeColumn(WriteBuffer & ss);
     void clear();
 
-    // FIXME: expose for ColumnArray::encodeIntoDatumData, find a better way to implement it.
-    void finishAppendVar(UInt32 size);
-    // WriteBufferFromOwnString is not moveable.
-    std::unique_ptr<WriteBufferFromOwnString> data;
-
 private:
+    void finishAppendVar(UInt32 size);
     void finishAppendFixed();
     bool isFixed() const { return fixed_size != VAR_SIZE; }
 
     void appendNullBitMap(bool value);
 
+    // WriteBufferFromOwnString is not moveable.
+    std::unique_ptr<WriteBufferFromOwnString> data;
+
     UInt32 length;
     UInt32 null_cnt;
     std::vector<UInt8> null_bitmap;

dbms/src/Flash/Coprocessor/tests/gtest_streaming_writer.cpp

@@ -41,7 +41,7 @@ class TestStreamingWriter : public testing::Test
     }
 
 public:
-    TestStreamingWriter() {}
+    TestStreamingWriter() = default;
 
     // Return 10 Int64 column.
     static std::vector<tipb::FieldType> makeFields()
@@ -93,7 +93,7 @@ struct MockStreamWriter
     {}
 
     void write(tipb::SelectResponse & response) { checker(response); }
-    bool isWritable() const { throw Exception("Unsupport async write"); }
+    static bool isWritable() { throw Exception("Unsupport async write"); }
 
 private:
     MockStreamWriterChecker checker;

dbms/src/TiDB/Schema/TiDB.h

@@ -56,6 +56,7 @@ using DB::Timestamp;
 // Column types.
 // In format:
 // TiDB type, int value, codec flag, CH type.
+// The int value is defined in pingcap/tidb pkg/parser/mysql/type.go
 #ifdef M
 #error "Please undefine macro M first."
 #endif