fix: cannot load kafka external streams created before sharding_expr feature (#351)

Author: Gimi Liang

src/Interpreters/InterpreterCreateQuery.cpp

@@ -903,20 +903,20 @@ void InterpreterCreateQuery::handleExternalStreamCreation(ASTCreateQuery & creat
 
     sharding_expr_field.tryGet<String>(sharding_expr);
 
-    ASTPtr sharding_expr_ast;
     if (sharding_expr.empty())
     {
-        sharding_expr_ast = makeASTFunction("rand");
+        create.storage->set(create.storage->engine, makeASTFunction("ExternalStream"));
     }
     else
     {
         ParserFunction parser;
         const char * begin{sharding_expr.data()};
         const char * end{begin + sharding_expr.size()};
-        sharding_expr_ast = parseQuery(parser, begin, end, "", 0, 0);
+        auto sharding_expr_ast = parseQuery(parser, begin, end, "", 0, 0);
+
+        create.storage->set(create.storage->engine, makeASTFunction("ExternalStream", sharding_expr_ast));
     }
 
-    create.storage->set(create.storage->engine, makeASTFunction("ExternalStream", sharding_expr_ast));
 
     if (create.storage->engine->name != "ExternalStream")
         throw Exception(ErrorCodes::INCORRECT_QUERY, "External stream requires ExternalStream engine");

src/Storages/ExternalStream/Kafka/Kafka.cpp

@@ -25,12 +25,12 @@ extern const int INVALID_SETTING_VALUE;
 extern const int RESOURCE_NOT_FOUND;
 }
 
-Kafka::Kafka(IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings_, ASTPtr partitioning_expr_, bool attach)
+Kafka::Kafka(IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings_, const ASTs & engine_args_, bool attach)
     : storage_id(storage->getStorageID())
     , settings(std::move(settings_))
     , data_format(settings->data_format.value)
     , log(&Poco::Logger::get("External-" + settings->topic.value))
-    , partitioning_expr(std::move(partitioning_expr_))
+    , engine_args(engine_args_)
 {
     assert(settings->type.value == StreamTypes::KAFKA || settings->type.value == StreamTypes::REDPANDA);
 

src/Storages/ExternalStream/Kafka/Kafka.h

@@ -13,7 +13,7 @@ class IStorage;
 class Kafka final : public StorageExternalStreamImpl
 {
 public:
-    Kafka(IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings_, ASTPtr partitioning_expr, bool attach);
+    Kafka(IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings_, const ASTs & engine_args_, bool attach);
     ~Kafka() override = default;
 
     void startup() override { }
@@ -42,7 +42,8 @@ class Kafka final : public StorageExternalStreamImpl
     const String & password() const { return settings->password.value; }
     const String & sslCaCertFile() const { return settings->ssl_ca_cert_file.value; }
     const klog::KConfParams & properties() const { return kafka_properties; }
-    const ASTPtr & partitioning_expr_ast() const { return partitioning_expr; }
+    bool hasCustomShardingExpr() const { return !engine_args.empty(); }
+    const ASTPtr & shardingExprAst() const { assert(!engine_args.empty()); return engine_args[0]; }
 
 private:
     void calculateDataFormat(const IStorage * storage);
@@ -61,7 +62,7 @@ class Kafka final : public StorageExternalStreamImpl
 
     NamesAndTypesList virtual_column_names_and_types;
     klog::KConfParams kafka_properties;
-    ASTPtr partitioning_expr;
+    const ASTs engine_args;
 
     std::mutex shards_mutex;
     int32_t shards = 0;

src/Storages/ExternalStream/Kafka/KafkaSink.cpp

@@ -24,55 +24,62 @@ extern const int TYPE_MISMATCH;
 
 namespace KafkaStream
 {
-ChunkPartitioner::ChunkPartitioner(ContextPtr context, const Block & header, const ASTPtr & partitioning_expr_ast)
+void ChunkSharder::useRandomSharding()
 {
-    /// `InterpreterCreateQuery::handleExternalStreamCreation` ensures this
-    assert(partitioning_expr_ast);
+    random_sharding = true;
+    std::random_device r;
+    rand = std::minstd_rand(r());
+}
+
+ChunkSharder::ChunkSharder(ContextPtr context, const Block & header, const ASTPtr & sharding_expr_ast)
+{
+    assert(sharding_expr_ast);
 
-    ASTPtr query = partitioning_expr_ast;
+    ASTPtr query = sharding_expr_ast;
     auto syntax_result = TreeRewriter(context).analyze(query, header.getNamesAndTypesList());
-    partitioning_expr = ExpressionAnalyzer(query, syntax_result, context).getActions(false);
+    sharding_expr = ExpressionAnalyzer(query, syntax_result, context).getActions(false);
 
-    partitioning_key_column_name = partitioning_expr_ast->getColumnName();
+    sharding_key_column_name = sharding_expr_ast->getColumnName();
 
-    if (auto * shard_func = partitioning_expr_ast->as<ASTFunction>())
+    if (auto * shard_func = sharding_expr_ast->as<ASTFunction>())
     {
         if (shard_func->name == "rand" || shard_func->name == "RAND")
-        {
-            random_partitioning = true;
-            std::random_device r;
-            rand = std::minstd_rand(r());
-        }
+            this->useRandomSharding();
     }
 }
 
-BlocksWithShard ChunkPartitioner::partition(Block block, Int32 partition_cnt) const
+ChunkSharder::ChunkSharder()
+{
+    this->useRandomSharding();
+}
+
+BlocksWithShard ChunkSharder::shard(Block block, Int32 shard_cnt) const
 {
     /// no topics have zero partitions
-    assert(partition_cnt > 0);
+    assert(shard_cnt > 0);
 
-    if (partition_cnt == 1)
+    if (shard_cnt == 1)
         return {BlockWithShard{Block(std::move(block)), 0}};
 
-    if (random_partitioning)
-        return {BlockWithShard{Block(std::move(block)), getNextShardIndex(partition_cnt)}};
+    if (random_sharding)
+        return {BlockWithShard{Block(std::move(block)), getNextShardIndex(shard_cnt)}};
 
-    return doParition(std::move(block), partition_cnt);
+    return doSharding(std::move(block), shard_cnt);
 }
 
-BlocksWithShard ChunkPartitioner::doParition(Block block, Int32 partition_cnt) const
+BlocksWithShard ChunkSharder::doSharding(Block block, Int32 shard_cnt) const
 {
-    auto selector = createSelector(block, partition_cnt);
+    auto selector = createSelector(block, shard_cnt);
 
-    Blocks partitioned_blocks{static_cast<size_t>(partition_cnt)};
+    Blocks partitioned_blocks{static_cast<size_t>(shard_cnt)};
 
-    for (Int32 i = 0; i < partition_cnt; ++i)
+    for (Int32 i = 0; i < shard_cnt; ++i)
         partitioned_blocks[i] = block.cloneEmpty();
 
     for (size_t pos = 0; pos < block.columns(); ++pos)
     {
-        MutableColumns partitioned_columns = block.getByPosition(pos).column->scatter(partition_cnt, selector);
-        for (Int32 i = 0; i < partition_cnt; ++i)
+        MutableColumns partitioned_columns = block.getByPosition(pos).column->scatter(shard_cnt, selector);
+        for (Int32 i = 0; i < shard_cnt; ++i)
             partitioned_blocks[i].getByPosition(pos).column = std::move(partitioned_columns[i]);
     }
 
@@ -89,14 +96,14 @@ BlocksWithShard ChunkPartitioner::doParition(Block block, Int32 partition_cnt) c
     return blocks_with_shard;
 }
 
-IColumn::Selector ChunkPartitioner::createSelector(Block block, Int32 partition_cnt) const
+IColumn::Selector ChunkSharder::createSelector(Block block, Int32 shard_cnt) const
 {
-    std::vector<UInt64> slot_to_shard(partition_cnt);
+    std::vector<UInt64> slot_to_shard(shard_cnt);
     std::iota(slot_to_shard.begin(), slot_to_shard.end(), 0);
 
-    partitioning_expr->execute(block);
+    sharding_expr->execute(block);
 
-    const auto & key_column = block.getByName(partitioning_key_column_name);
+    const auto & key_column = block.getByName(sharding_key_column_name);
 
 /// If key_column.type is DataTypeLowCardinality, do shard according to its dictionaryType
 #define CREATE_FOR_TYPE(TYPE) \
@@ -236,7 +243,10 @@ KafkaSink::KafkaSink(const Kafka * kafka, const Block & header, ContextPtr conte
         writer = FormatFactory::instance().getOutputFormat(data_format, *wb, header, context);
     writer->setAutoFlush();
 
-    partitioner = std::make_unique<KafkaStream::ChunkPartitioner>(context, header, kafka->partitioning_expr_ast());
+    if (kafka->hasCustomShardingExpr())
+        partitioner = std::make_unique<KafkaStream::ChunkSharder>(context, header, kafka->shardingExprAst());
+    else
+        partitioner = std::make_unique<KafkaStream::ChunkSharder>();
 
     polling_threads.scheduleOrThrowOnError([this]() {
         while (!is_finished.test())
@@ -251,7 +261,7 @@ void KafkaSink::consume(Chunk chunk)
         return;
 
     auto block = getHeader().cloneWithColumns(chunk.detachColumns());
-    auto blocks = partitioner->partition(std::move(block), partition_cnt);
+    auto blocks = partitioner->shard(std::move(block), partition_cnt);
 
     for (auto & blockWithShard : blocks)
     {

src/Storages/ExternalStream/Kafka/KafkaSink.h

@@ -19,23 +19,26 @@ namespace DB
 
 namespace KafkaStream
 {
-/// Partition Chunk's to paritions (i.e. shards) by the partitioning expression.
-class ChunkPartitioner
+/// Shard Chunk's to shards (or partitions in Kafka's term) by the sharding expression.
+class ChunkSharder
 {
 public:
-    ChunkPartitioner(ContextPtr context, const Block & header, const ASTPtr & partitioning_expr_ast);
-    BlocksWithShard partition(Block block, Int32 partition_cnt) const;
+    ChunkSharder(ContextPtr context, const Block & header, const ASTPtr & sharding_expr_ast);
+    ChunkSharder();
+
+    BlocksWithShard shard(Block block, Int32 shard_cnt) const;
 
 private:
-    Int32 getNextShardIndex(Int32 partition_cnt) const noexcept { return static_cast<Int32>(rand()) % partition_cnt; }
+    void useRandomSharding();
+    Int32 getNextShardIndex(Int32 shard_cnt) const noexcept { return static_cast<Int32>(rand()) % shard_cnt; }
 
-    BlocksWithShard doParition(Block block, Int32 partition_cnt) const;
+    BlocksWithShard doSharding(Block block, Int32 shard_cnt) const;
 
-    IColumn::Selector createSelector(Block block, Int32 partition_cnt) const;
+    IColumn::Selector createSelector(Block block, Int32 shard_cnt) const;
 
-    ExpressionActionsPtr partitioning_expr;
-    String partitioning_key_column_name;
-    bool random_partitioning = false;
+    ExpressionActionsPtr sharding_expr;
+    String sharding_key_column_name;
+    bool random_sharding = false;
     mutable std::minstd_rand rand;
 };
 }
@@ -87,7 +90,7 @@ class KafkaSink final : public SinkToStorage
     ThreadPool polling_threads;
     std::atomic_flag is_finished;
     Int32 partition_cnt;
-    std::unique_ptr<KafkaStream::ChunkPartitioner> partitioner;
+    std::unique_ptr<KafkaStream::ChunkSharder> partitioner;
 
     Poco::Logger * log;
 };

src/Storages/ExternalStream/StorageExternalStream.cpp

@@ -39,8 +39,8 @@ buildShardingKeyExpression(ASTPtr sharding_key, ContextPtr context, const NamesA
 }
 
 void validateEngineArgs(ContextPtr context, ASTs & engine_args, const ColumnsDescription & columns) {
-    if (engine_args.size() == 0)
-        throw Exception(ErrorCodes::BAD_ARGUMENTS, "ExternalStream requires sharding key expression");
+    if (engine_args.empty())
+        return;
 
     auto sharding_expr = buildShardingKeyExpression(engine_args[0], context, columns.getAllPhysical());
     const auto & block = sharding_expr->getSampleBlock();
@@ -55,13 +55,13 @@ void validateEngineArgs(ContextPtr context, ASTs & engine_args, const ColumnsDes
 }
 
 std::unique_ptr<StorageExternalStreamImpl> createExternalStream(
-    IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings, ContextPtr & context [[maybe_unused]], const ASTPtr & sharding_expr, bool attach)
+    IStorage * storage, std::unique_ptr<ExternalStreamSettings> settings, ContextPtr & context [[maybe_unused]], const ASTs & engine_args, bool attach)
 {
     if (settings->type.value.empty())
         throw Exception(ErrorCodes::BAD_ARGUMENTS, "External stream type is required in settings");
 
     if (settings->type.value == StreamTypes::KAFKA || settings->type.value == StreamTypes::REDPANDA)
-        return std::make_unique<Kafka>(storage, std::move(settings), sharding_expr, attach);
+        return std::make_unique<Kafka>(storage, std::move(settings), engine_args, attach);
 #ifdef OS_LINUX
     else if (settings->type.value == StreamTypes::LOG && context->getSettingsRef()._tp_enable_log_stream_expr.value)
         return std::make_unique<FileLog>(storage, std::move(settings));
@@ -130,7 +130,7 @@ SinkToStoragePtr StorageExternalStream::write(const ASTPtr & query, const Storag
 }
 
 StorageExternalStream::StorageExternalStream(
-    const ASTPtr & sharding_key_,
+    const ASTs & engine_args,
     const StorageID & table_id_,
     ContextPtr context_,
     const ColumnsDescription & columns_,
@@ -143,7 +143,7 @@ StorageExternalStream::StorageExternalStream(
     setInMemoryMetadata(storage_metadata);
 
 
-    auto stream = createExternalStream(this, std::move(external_stream_settings_), context_, sharding_key_, attach);
+    auto stream = createExternalStream(this, std::move(external_stream_settings_), context_, engine_args, attach);
     external_stream.swap(stream);
 }
 
@@ -161,7 +161,7 @@ void registerStorageExternalStream(StorageFactory & factory)
             external_stream_settings->loadFromQuery(*args.storage_def);
 
             return StorageExternalStream::create(
-                args.engine_args[0], args.table_id, args.getContext(), args.columns, std::move(external_stream_settings), args.attach);
+                args.engine_args, args.table_id, args.getContext(), args.columns, std::move(external_stream_settings), args.attach);
         }
         else
             throw Exception(ErrorCodes::BAD_ARGUMENTS, "External stream requires correct settings setup");

src/Storages/ExternalStream/StorageExternalStream.h

@@ -52,7 +52,7 @@ class StorageExternalStream final : public shared_ptr_helper<StorageExternalStre
 
 protected:
     StorageExternalStream(
-        const ASTPtr & sharding_key_,
+        const ASTs & engine_args,
         const StorageID & table_id_,
         ContextPtr context_,
         const ColumnsDescription & columns_,