[Vectorized][Feature] Support mysql external table insert into stm

be/src/exec/data_sink.cpp

@@ -37,6 +37,7 @@
 
 #include "vec/sink/result_sink.h"
 #include "vec/sink/vdata_stream_sender.h"
+#include "vec/sink/vmysql_table_writer.h"
 #include "vec/sink/vtablet_sink.h"
 
 namespace doris {
@@ -113,10 +114,14 @@ Status DataSink::create_data_sink(ObjectPool* pool, const TDataSink& thrift_sink
         if (!thrift_sink.__isset.mysql_table_sink) {
             return Status::InternalError("Missing data buffer sink.");
         }
-
-        // TODO: figure out good buffer size based on size of output row
-        MysqlTableSink* mysql_tbl_sink = new MysqlTableSink(pool, row_desc, output_exprs);
-        sink->reset(mysql_tbl_sink);
+        if (is_vec) {
+            doris::vectorized::VMysqlTableSink* vmysql_tbl_sink = new doris::vectorized::VMysqlTableSink(pool, row_desc, output_exprs);
+            sink->reset(vmysql_tbl_sink);
+        } else {
+            // TODO: figure out good buffer size based on size of output row
+            MysqlTableSink* mysql_tbl_sink = new MysqlTableSink(pool, row_desc, output_exprs);
+            sink->reset(mysql_tbl_sink);
+        }
         break;
 #else
         return Status::InternalError(

be/src/runtime/mysql_table_writer.cpp

@@ -26,6 +26,12 @@
 #include "runtime/row_batch.h"
 #include "runtime/tuple_row.h"
 #include "util/types.h"
+#include "vec/columns/column_nullable.h"
+#include "vec/core/block.h"
+#include "vec/data_types/data_type.h"
+#include "vec/exprs/vexpr.h"
+#include "vec/exprs/vexpr_context.h"
+#include "vec/core/materialize_block.h"
 
 namespace doris {
 
@@ -38,7 +44,10 @@ std::string MysqlConnInfo::debug_string() const {
 }
 
 MysqlTableWriter::MysqlTableWriter(const std::vector<ExprContext*>& output_expr_ctxs)
-        : _output_expr_ctxs(output_expr_ctxs) {}
+        : _output_expr_ctxs(output_expr_ctxs), _vec_output_expr_ctxs {} {}
+
+MysqlTableWriter::MysqlTableWriter(const std::vector<vectorized::VExprContext*>& output_expr_ctxs)
+        : _output_expr_ctxs {}, _vec_output_expr_ctxs(output_expr_ctxs) {}
 
 MysqlTableWriter::~MysqlTableWriter() {
     if (_mysql_conn) {
@@ -182,4 +191,141 @@ Status MysqlTableWriter::append(RowBatch* batch) {
     return Status::OK();
 }
 
+Status MysqlTableWriter::append(vectorized::Block* block) {
+    Status status = Status::OK();
+    if (block == nullptr || block->rows() == 0) {
+        return status;
+    }
+
+    auto output_block = vectorized::VExprContext::get_output_block_after_execute_exprs(
+            _vec_output_expr_ctxs, *block, status);
+
+    auto num_rows = output_block.rows();
+    if (UNLIKELY(num_rows == 0)) {
+        return status;
+    }
+    materialize_block_inplace(output_block);
+    for (int i = 0; i < num_rows; ++i) {
+        RETURN_IF_ERROR(insert_row(output_block, i));
+    }
+    return Status::OK();
+}
+
+Status MysqlTableWriter::insert_row(vectorized::Block& block, size_t row) {
+    _insert_stmt_buffer.clear();
+    fmt::format_to(_insert_stmt_buffer, "INSERT INTO {} VALUES (", _mysql_tbl);
+    int num_columns = _vec_output_expr_ctxs.size();
+
+    for (int i = 0; i < num_columns; ++i) {
+        auto& column_ptr = block.get_by_position(i).column;
+        auto& type_ptr = block.get_by_position(i).type;
+
+        if (i != 0) {
+            fmt::format_to(_insert_stmt_buffer, "{}", ", ");
+        }
+
+        vectorized::ColumnPtr column;
+        if (type_ptr->is_nullable()) {
+            column = assert_cast<const vectorized::ColumnNullable&>(*column_ptr)
+                             .get_nested_column_ptr();
+            if (column_ptr->is_null_at(row)) {
+                fmt::format_to(_insert_stmt_buffer, "{}", "NULL");
+                continue;
+            }
+        } else {
+            column = column_ptr;
+        }
+
+        switch (_vec_output_expr_ctxs[i]->root()->result_type()) {
+        case TYPE_BOOLEAN: {
+            auto& data = assert_cast<const vectorized::ColumnUInt8&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_TINYINT: {
+            auto& data = assert_cast<const vectorized::ColumnInt8&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_SMALLINT: {
+            auto& data = assert_cast<const vectorized::ColumnInt16&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_INT: {
+            auto& data = assert_cast<const vectorized::ColumnInt32&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_BIGINT: {
+            auto& data = assert_cast<const vectorized::ColumnInt64&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_FLOAT: {
+            auto& data = assert_cast<const vectorized::ColumnFloat32&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+        case TYPE_DOUBLE: {
+            auto& data = assert_cast<const vectorized::ColumnFloat64&>(*column).get_data();
+            fmt::format_to(_insert_stmt_buffer, "{}", data[row]);
+            break;
+        }
+
+        case TYPE_STRING:
+        case TYPE_CHAR:
+        case TYPE_VARCHAR: {
+            const auto& string_val =
+                    assert_cast<const vectorized::ColumnString&>(*column).get_data_at(row);
+            DCHECK(string_val.data != nullptr);
+            std::unique_ptr<char[]> buf(new char[2 * string_val.size + 1]);
+            mysql_real_escape_string(_mysql_conn, buf.get(), string_val.data, string_val.size);
+            fmt::format_to(_insert_stmt_buffer, "'{}'", buf.get());
+
+            break;
+        }
+        case TYPE_DECIMALV2: {
+            DecimalV2Value value =
+                    (DecimalV2Value)
+                            assert_cast<const vectorized::ColumnDecimal<vectorized::Decimal128>&>(
+                                    *column)
+                                    .get_data()[row];
+            fmt::format_to(_insert_stmt_buffer, "{}", value.to_string());
+            break;
+        }
+        case TYPE_DATE:
+        case TYPE_DATETIME: {
+            int64_t int_val = assert_cast<const vectorized::ColumnInt64&>(*column).get_data()[row];
+            vectorized::VecDateTimeValue value =
+                    binary_cast<int64_t, doris::vectorized::VecDateTimeValue>(int_val);
+
+            char buf[64];
+            char* pos = value.to_string(buf);
+            std::string str(buf, pos - buf - 1);
+            fmt::format_to(_insert_stmt_buffer, "'{}'", str);
+            break;
+        }
+        default: {
+            fmt::memory_buffer err_out;
+            fmt::format_to(err_out, "can't convert this type to mysql type. type = {}",
+                           _vec_output_expr_ctxs[i]->root()->type().type);
+            return Status::InternalError(err_out.data());
+        }
+        }
+    }
+
+    fmt::format_to(_insert_stmt_buffer, "{}", ")");
+
+    // Insert this to MySQL server
+    if (mysql_real_query(_mysql_conn, _insert_stmt_buffer.data(), _insert_stmt_buffer.size())) {
+        fmt::memory_buffer err_ss;
+        fmt::format_to(err_ss, "Insert to mysql server({}) failed, because: {}.",
+                       mysql_get_host_info(_mysql_conn), mysql_error(_mysql_conn));
+        return Status::InternalError(err_ss.data());
+    }
+
+    return Status::OK();
+}
+
 } // namespace doris

be/src/runtime/mysql_table_writer.h

@@ -18,11 +18,11 @@
 #ifndef DORIS_BE_RUNTIME_MYSQL_TABLE_WRITER_H
 #define DORIS_BE_RUNTIME_MYSQL_TABLE_WRITER_H
 
+#include <fmt/format.h>
 #include <string>
 #include <vector>
 
 #include "common/status.h"
-
 #ifndef __DorisMysql
 #define __DorisMysql void
 #endif
@@ -42,10 +42,15 @@ struct MysqlConnInfo {
 class RowBatch;
 class TupleRow;
 class ExprContext;
+namespace vectorized {
+class VExprContext;
+class Block;
+}
 
 class MysqlTableWriter {
 public:
     MysqlTableWriter(const std::vector<ExprContext*>& output_exprs);
+    MysqlTableWriter(const std::vector<vectorized::VExprContext*>& output_exprs);
     ~MysqlTableWriter();
 
     // connect to mysql server
@@ -55,6 +60,8 @@ class MysqlTableWriter {
 
     Status append(RowBatch* batch);
 
+    Status append(vectorized::Block* block);
+
     Status abort_tarns() { return Status::OK(); }
 
     Status finish_tarns() { return Status::OK(); }
@@ -63,6 +70,10 @@ class MysqlTableWriter {
     Status insert_row(TupleRow* row);
 
     const std::vector<ExprContext*>& _output_expr_ctxs;
+    //vectorized mode insert_row
+    Status insert_row(vectorized::Block& block, size_t row);
+    const std::vector<vectorized::VExprContext*>& _vec_output_expr_ctxs;
+    fmt::memory_buffer _insert_stmt_buffer;
     std::string _mysql_tbl;
     __DorisMysql* _mysql_conn;
 };

be/src/vec/CMakeLists.txt

@@ -149,6 +149,7 @@ set(VEC_FILES
   sink/result_sink.cpp
   sink/vdata_stream_sender.cpp
   sink/vtablet_sink.cpp
+  sink/vmysql_table_writer.cpp
   runtime/vdatetime_value.cpp
   runtime/vdata_stream_recvr.cpp
   runtime/vdata_stream_mgr.cpp

be/src/vec/sink/vmysql_table_writer.cpp

@@ -0,0 +1,91 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <sstream>
+
+#include "exprs/expr.h"
+#include "runtime/mem_tracker.h"
+#include "runtime/mysql_table_sink.h"
+#include "runtime/runtime_state.h"
+#include "util/debug_util.h"
+#include "util/runtime_profile.h"
+#include "vec/sink/vmysql_table_writer.h"
+#include "vec/exprs/vexpr.h"
+
+namespace doris {
+namespace vectorized {
+VMysqlTableSink::VMysqlTableSink(ObjectPool* pool, const RowDescriptor& row_desc,
+                               const std::vector<TExpr>& t_exprs)
+        : _pool(pool),
+          _row_desc(row_desc),
+          _t_output_expr(t_exprs),
+          _mem_tracker(MemTracker::CreateTracker(-1, "VMysqlTableSink")) {
+    _name = "VMysqlTableSink";
+}
+
+VMysqlTableSink::~VMysqlTableSink() {}
+
+Status VMysqlTableSink::init(const TDataSink& t_sink) {
+    RETURN_IF_ERROR(DataSink::init(t_sink));
+    const TMysqlTableSink& t_mysql_sink = t_sink.mysql_table_sink;
+
+    _conn_info.host = t_mysql_sink.host;
+    _conn_info.port = t_mysql_sink.port;
+    _conn_info.user = t_mysql_sink.user;
+    _conn_info.passwd = t_mysql_sink.passwd;
+    _conn_info.db = t_mysql_sink.db;
+    _mysql_tbl = t_mysql_sink.table;
+
+    // From the thrift expressions create the real exprs.
+    RETURN_IF_ERROR(VExpr::create_expr_trees(_pool, _t_output_expr, &_output_expr_ctxs));
+    return Status::OK();
+}
+
+Status VMysqlTableSink::prepare(RuntimeState* state) {
+    RETURN_IF_ERROR(DataSink::prepare(state));
+    // Prepare the exprs to run.
+    RETURN_IF_ERROR(VExpr::prepare(_output_expr_ctxs, state, _row_desc, _mem_tracker));
+    std::stringstream title;
+    title << "VMysqlTableSink (frag_id=" << state->fragment_instance_id() << ")";
+    // create profile
+    _profile = state->obj_pool()->add(new RuntimeProfile(title.str()));
+    return Status::OK();
+}
+
+Status VMysqlTableSink::open(RuntimeState* state) {
+    // Prepare the exprs to run.
+    RETURN_IF_ERROR(VExpr::open(_output_expr_ctxs, state));
+    // create writer
+    _writer = state->obj_pool()->add(new doris::MysqlTableWriter(_output_expr_ctxs));
+    RETURN_IF_ERROR(_writer->open(_conn_info, _mysql_tbl));
+    return Status::OK();
+}
+
+Status VMysqlTableSink::send(RuntimeState* state, RowBatch* batch) {
+    return Status::NotSupported("Not Implemented VMysqlTableSink::send(RuntimeState* state, RowBatch* batch)");
+}
+
+Status VMysqlTableSink::send(RuntimeState* state, Block* block) {
+    return _writer->append(block);
+}
+
+Status VMysqlTableSink::close(RuntimeState* state, Status exec_status) {
+    VExpr::close(_output_expr_ctxs, state);
+    return Status::OK();
+}
+} // namespace vectorized
+} // namespace doris