[Kafka External Stream] bugfix: repect data format for count optimiza…

…tion (#762)

src/Storages/ExternalStream/Kafka/Kafka.cpp

@@ -259,6 +259,17 @@ Kafka::Kafka(
 
     calculateDataFormat(storage);
 
+    /// For raw format, one_message_per_row default to true. Because if multiple rows are combined together,
+    /// it could be hard to separate them again.
+    if (data_format == "RawBLOB" && !settings->isChanged("one_message_per_row"))
+        settings->set("one_message_per_row", true);
+
+    /// Only optimize trivial count on some specific formats, and settings.
+    support_count_optimization = data_format == "RawBLOB" ||
+        data_format == "ProtobufSingle" ||
+        (data_format == "Avro" && (!settings->format_schema.value.empty() || !settings->kafka_schema_registry_url.value.empty())) ||
+        settings->one_message_per_row.value;
+
     cacheVirtualColumnNamesAndTypes();
 
     rd_kafka_conf_set_log_cb(conf.get(), &Kafka::onLog);
@@ -447,6 +458,9 @@ std::vector<Int32> getShardsToQuery(const String & shards_exp, Int32 parition_co
 
 std::optional<UInt64> Kafka::totalRows(const Settings & settings_ref)
 {
+    if (!support_count_optimization)
+        return {};
+
     auto consumer = getConsumer();
     auto topic = RdKafka::Topic(*consumer->getHandle(), topicName());
     auto shards_to_query = getShardsToQuery(settings_ref.shards.value, topic.getPartitionCount());

src/Storages/ExternalStream/Kafka/Kafka.h

@@ -111,6 +111,8 @@ class Kafka final : public StorageExternalStreamImpl
     std::vector<Int32> shards_from_settings;
     fs::path broker_ca_file;
 
+    bool support_count_optimization = false;
+
     ConfPtr conf;
     UInt64 poll_timeout_ms = 0;