enhancement: properly produce raw format data to Kafka external streams (#436)

Author: zliang-min

src/Storages/ExternalStream/Kafka/Kafka.cpp

@@ -183,24 +183,19 @@ void Kafka::calculateDataFormat(const IStorage * storage)
     if (!data_format.empty())
         return;
 
+    /// If there is only one column and its type is a string type, use RawBLOB. Use JSONEachRow otherwise.
     auto column_names_and_types{storage->getInMemoryMetadata().getColumns().getOrdinary()};
-    if (column_names_and_types.size() != 1)
-        throw Exception(
-            ErrorCodes::INVALID_SETTING_VALUE,
-            "`data_format` settings is empty but the Kafka external stream definition has multiple columns. Proton doesn't know how to "
-            "parse Kafka messages without a data format.");
-
-    auto type = column_names_and_types.begin()->type;
-    if (type->getTypeId() == TypeIndex::String || type->getTypeId() == TypeIndex::FixedString)
+    if (column_names_and_types.size() == 1)
     {
-        /// no-op
+        auto type = column_names_and_types.begin()->type->getTypeId();
+        if (type == TypeIndex::String || type == TypeIndex::FixedString)
+        {
+            data_format = "RawBLOB";
+            return;
+        }
     }
-    /// FIXME: JSON logic
-    else if (type->getTypeId() == TypeIndex::Object)
-        data_format = "JSONEachRow";
-    else
-        throw Exception(
-            ErrorCodes::NOT_IMPLEMENTED, "Automatically converting Kafka message to {} type is not supported yet", type->getName());
+
+    data_format = "JSONEachRow";
 }
 
 /// FIXME, refactor out as util and unit test it

src/Storages/ExternalStream/Kafka/KafkaSink.cpp

@@ -221,9 +221,8 @@ KafkaSink::KafkaSink(const Kafka * kafka, const Block & header, ContextPtr conte
     topic = klog::KTopicPtr(rd_kafka_topic_new(producer.get(), kafka->topic().c_str(), nullptr), rd_kafka_topic_destroy);
     wb = std::make_unique<WriteBufferFromKafka>(topic.get());
 
-    String data_format = kafka->dataFormat();
-    if (data_format.empty())
-        data_format = "JSONEachRow";
+    const auto & data_format = kafka->dataFormat();
+    assert(!data_format.empty());
 
     /// The callback allows `IRowOutputFormat` based formats produce one Kafka message per row.
     if (kafka->produceOneMessagePerRow())

src/Storages/ExternalStream/Kafka/KafkaSource.cpp

@@ -130,62 +130,10 @@ void KafkaSource::parseMessage(void * kmessage, size_t total_count, void * data)
 
 void KafkaSource::doParseMessage(const rd_kafka_message_t * kmessage, size_t /*total_count*/)
 {
-    if (format_executor)
-        parseFormat(kmessage);
-    else
-        parseRaw(kmessage);
-
+    parseFormat(kmessage);
     ckpt_data.last_sn = kmessage->offset;
 }
 
-void KafkaSource::parseRaw(const rd_kafka_message_t * kmessage)
-{
-    if (!request_virtual_columns)
-    {
-        /// fast path
-        assert(physical_header.columns() == 1);
-
-        if (current_batch.empty())
-            current_batch.push_back(physical_header.getByPosition(0).type->createColumn());
-
-        current_batch.back()->insertData(static_cast<const char *>(kmessage->payload), kmessage->len);
-        external_stream_counter->addToReadBytes(kmessage->len);
-        external_stream_counter->addToReadCounts(1);
-    }
-    else
-    {
-        /// slower path, request virtual columns
-        if (!current_batch.empty())
-        {
-            assert(current_batch.size() == virtual_col_value_functions.size());
-            for (size_t i = 0, n = virtual_col_value_functions.size(); i < n; ++i)
-            {
-                if (!virtual_col_value_functions[i])
-                    current_batch[i]->insertData(static_cast<const char *>(kmessage->payload), kmessage->len);
-                else
-                    current_batch[i]->insertMany(virtual_col_value_functions[i](kmessage), 1);
-            }
-        }
-        else
-        {
-            for (size_t i = 0, n = virtual_col_value_functions.size(); i < n; ++i)
-            {
-                if (!virtual_col_value_functions[i])
-                {
-                    current_batch.push_back(physical_header.getByPosition(0).type->createColumn());
-                    current_batch.back()->insertData(static_cast<const char *>(kmessage->payload), kmessage->len);
-                }
-                else
-                {
-                    auto column = virtual_col_types[i]->createColumn();
-                    column->insertMany(virtual_col_value_functions[i](kmessage), 1);
-                    current_batch.push_back(std::move(column));
-                }
-            }
-        }
-    }
-}
-
 void KafkaSource::parseFormat(const rd_kafka_message_t * kmessage)
 {
     assert(format_executor);
@@ -285,21 +233,19 @@ void KafkaSource::initConsumer(const Kafka * kafka)
 void KafkaSource::initFormatExecutor(const Kafka * kafka)
 {
     const auto & data_format = kafka->dataFormat();
-    if (!data_format.empty())
-    {
-        auto input_format
-            = FormatFactory::instance().getInputFormat(data_format, read_buffer, non_virtual_header, query_context, max_block_size, kafka->getFormatSettings(query_context));
 
-        format_executor = std::make_unique<StreamingFormatExecutor>(
-            non_virtual_header, std::move(input_format), [](const MutableColumns &, Exception &) -> size_t { return 0; });
+    auto input_format
+        = FormatFactory::instance().getInputFormat(data_format, read_buffer, non_virtual_header, query_context, max_block_size);
 
-        auto converting_dag = ActionsDAG::makeConvertingActions(
-            non_virtual_header.cloneEmpty().getColumnsWithTypeAndName(),
-            physical_header.cloneEmpty().getColumnsWithTypeAndName(),
-            ActionsDAG::MatchColumnsMode::Name);
+    format_executor = std::make_unique<StreamingFormatExecutor>(
+        non_virtual_header, std::move(input_format), [](const MutableColumns &, Exception &) -> size_t { return 0; });
 
-        convert_non_virtual_to_physical_action = std::make_shared<ExpressionActions>(std::move(converting_dag));
-    }
+    auto converting_dag = ActionsDAG::makeConvertingActions(
+        non_virtual_header.cloneEmpty().getColumnsWithTypeAndName(),
+        physical_header.cloneEmpty().getColumnsWithTypeAndName(),
+        ActionsDAG::MatchColumnsMode::Name);
+
+    convert_non_virtual_to_physical_action = std::make_shared<ExpressionActions>(std::move(converting_dag));
 }
 
 void KafkaSource::calculateColumnPositions()

src/Storages/ExternalStream/Kafka/KafkaSource.h

@@ -53,7 +53,6 @@ class KafkaSource final : public ISource
     static void parseMessage(void * kmessage, size_t total_count, void * data);
     void doParseMessage(const rd_kafka_message_s * kmessage, size_t total_count);
     void parseFormat(const rd_kafka_message_s * kmessage);
-    void parseRaw(const rd_kafka_message_s * kmessage);
 
     inline void readAndProcess();