support monotonic function and support _tp_sn, _tp_shard in random s… (#359)

lizhou1111 · web-flow · commit cc67ce08b00a · 2023-12-05T14:09:21.000+08:00
diff --git a/src/Functions/FunctionMonotonic.cpp b/src/Functions/FunctionMonotonic.cpp
@@ -0,0 +1,113 @@
+#include <Columns/ColumnConst.h>
+#include <Core/Field.h>
+#include <DataTypes/DataTypeFactory.h>
+#include <DataTypes/DataTypesNumber.h>
+#include <Formats/FormatSettings.h>
+#include <Functions/FunctionFactory.h>
+#include <Functions/IFunction.h>
+#include <IO/ReadHelpers.h>
+#include <IO/WriteHelpers.h>
+#include <base/map.h>
+
+#include <numeric>
+
+
+namespace DB
+{
+namespace ErrorCodes
+{
+extern const int BAD_ARGUMENTS;
+extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
+}
+
+
+/**
+ * Function monotonic() returns a monotonically increasing sequence of numbers.
+ * Now the argument range is min(Int64) ~ max(Int64).
+ * So it supports passing uint8, uint16, uint32, int8, int16, int32, int64 as argument.
+*/
+class FuntionMonotonic : public IFunction
+{
+public:
+    explicit FuntionMonotonic(Int64 start_num_) : start_num(start_num_) { }
+
+    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
+
+    size_t getNumberOfArguments() const override { return 1; }
+
+    bool isStateful() const override { return true; }
+
+    bool isDeterministic() const override { return false; }
+
+    bool isDeterministicInScopeOfQuery() const override { return false; }
+
+    bool isSuitableForConstantFolding() const override { return false; }
+
+    ColumnPtr executeImpl(const ColumnsWithTypeAndName &, const DataTypePtr & result_type, size_t input_rows_count) const override
+    {
+        auto column = ColumnInt64::create(input_rows_count);
+        std::iota(column->getData().begin(), column->getData().end(), start_num);
+        start_num += input_rows_count;
+        return column;
+    }
+
+    void serialize(WriteBuffer & wb) const override { writeIntBinary<Int64>(start_num, wb); }
+
+    void deserialize(ReadBuffer & rb) const override { readIntBinary<Int64>(start_num, rb); }
+
+    String getName() const override { return "monotonic"; }
+
+private:
+    mutable Int64 start_num;
+};
+
+class MonotonicOverloadResolver : public IFunctionOverloadResolver
+{
+public:
+    static constexpr auto name = "monotonic";
+
+    String getName() const override { return name; }
+
+    static FunctionOverloadResolverPtr create(ContextPtr) { return std::make_unique<MonotonicOverloadResolver>(); }
+
+    size_t getNumberOfArguments() const override { return 1; }
+
+    bool isStateful() const override { return true; }
+
+    bool isVariadic() const override { return false; }
+
+    bool isDeterministic() const override { return false; }
+
+    bool isDeterministicInScopeOfQuery() const override { return false; }
+    DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
+    {
+        if (arguments.size() != 1)
+            throw Exception("Function " + getName() + " requires 1 arguments", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+
+        if (!isInteger(arguments[0].type))
+            throw Exception(
+                fmt::format("Function {} requires integer argumengt but given {}", getName(), arguments[0].type->getFamilyName()),
+                ErrorCodes::BAD_ARGUMENTS);
+
+        return std::make_unique<DataTypeInt64>();
+    }
+
+    FunctionBasePtr buildImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & return_type) const override
+    {
+        if (!isColumnConst(*arguments[0].column))
+            throw Exception("Function " + getName() + " requires constant argument", ErrorCodes::BAD_ARGUMENTS);
+
+        auto start_num = assert_cast<const ColumnConst &>(*arguments[0].column).getInt(0);
+        return std::make_unique<FunctionToFunctionBaseAdaptor>(
+            std::make_unique<FuntionMonotonic>(start_num),
+            collections::map<DataTypes>(arguments, [](const auto & elem) { return elem.type; }),
+            return_type);
+    }
+};
+
+
+REGISTER_FUNCTION(Monotonic)
+{
+    factory.registerFunction<MonotonicOverloadResolver>({}, FunctionFactory::CaseInsensitive);
+}
+}
diff --git a/src/Storages/Streaming/StorageRandom.cpp b/src/Storages/Streaming/StorageRandom.cpp
@@ -36,7 +36,6 @@
 #include <Common/logger_useful.h>
 #include <Common/randomSeed.h>
 
-
 namespace DB
 {
 
@@ -66,7 +65,6 @@ void fillBufferWithRandomData(char * __restrict data, size_t size, pcg64 & rng)
     }
 }
 
-
 ColumnPtr fillColumnWithRandomData(const DataTypePtr type, UInt64 limit, pcg64 & rng, ContextPtr context)
 {
     TypeIndex idx = type->getTypeId();
@@ -353,6 +351,28 @@ ColumnPtr fillColumnWithRandomData(const DataTypePtr type, UInt64 limit, pcg64 &
     }
 }
 
+ColumnPtr
+fillColumnWithData(const DataTypePtr type, UInt64 limit, std::tuple<Int64, Int32, pcg64> & data, ContextPtr context, String col_name)
+{
+    if (col_name == ProtonConsts::RESERVED_SHARD)
+    {
+        auto & shard_num = std::get<0>(data);
+        auto column = type->createColumnConst(limit, shard_num)->convertToFullColumnIfConst();
+        return column;
+    }
+    else if (col_name == ProtonConsts::RESERVED_EVENT_SEQUENCE_ID)
+    {
+        auto & sn = std::get<1>(data);
+        auto column = type->createColumnConst(limit, sn)->convertToFullColumnIfConst();
+        sn++;
+        return column;
+    }
+    else
+    {
+        auto & rng = std::get<2>(data);
+        return fillColumnWithRandomData(type, limit, rng, context);
+    }
+}
 
 class GenerateRandomSource final : public ISource
 {
@@ -366,12 +386,12 @@ class GenerateRandomSource final : public ISource
         UInt64 events_per_second_,
         UInt64 interval_time_,
         bool is_streaming_,
-        UInt64 total_events_)
+        UInt64 total_events_,
+        size_t shard_num_)
         : ISource(Nested::flatten(prepareBlockToFill(block_header_)), true, ProcessorID::GenerateRandomSourceID)
         , block_size(block_size_)
         , block_full(std::move(block_header_))
         , our_columns(our_columns_)
-        , rng(random_seed_)
         , context(context_)
         , events_per_second(events_per_second_)
         , header_chunk(Nested::flatten(block_full.cloneEmpty()).getColumns(), 0)
@@ -380,7 +400,7 @@ class GenerateRandomSource final : public ISource
         , log(&Poco::Logger::get("GenerateRandSource"))
     {
         is_streaming = is_streaming_;
-
+        data_generate_helper = std::make_tuple(shard_num_, 1, pcg64(random_seed_));
         if (total_events == 0 && !is_streaming)
             total_events = events_per_second ? events_per_second : block_size;
 
@@ -414,8 +434,17 @@ class GenerateRandomSource final : public ISource
                 != ProtonConsts::RESERVED_COLUMN_NAMES.end();
             if (is_reserved_column || our_columns.hasDefault(elem.name))
                 continue;
+
             block_to_fill.insert(elem);
         }
+
+        auto dag
+            = evaluateMissingDefaults(block_to_fill, block_full.getNamesAndTypesList(), our_columns, context, true, false, true);
+        if (dag)
+        {
+            default_actions = std::make_shared<ExpressionActions>(
+                std::move(dag), ExpressionActionsSettings::fromContext(context, CompileExpressions::yes));
+        }
     }
 
     String getName() const override { return "Random"; }
@@ -493,7 +522,7 @@ class GenerateRandomSource final : public ISource
         Block block_to_fill_as_result(block_to_fill.cloneEmpty());
 
         for (const auto & elem : block_to_fill_as_result)
-            columns.emplace_back(fillColumnWithRandomData(elem.type, block_size_, rng, context));
+            columns.emplace_back(fillColumnWithData(elem.type, block_size_, data_generate_helper, context, elem.name));
 
         block_to_fill_as_result.setColumns(columns);
 
@@ -504,13 +533,8 @@ class GenerateRandomSource final : public ISource
             block_to_fill_as_result.insert(
                 {ColumnConst::create(ColumnUInt8::create(1, 0), block_size_), std::make_shared<DataTypeUInt8>(), "_dummy"});
 
-        auto dag = evaluateMissingDefaults(block_to_fill_as_result, block_full.getNamesAndTypesList(), our_columns, context);
-        if (dag)
-        {
-            auto actions = std::make_shared<ExpressionActions>(
-                std::move(dag), ExpressionActionsSettings::fromContext(context, CompileExpressions::yes));
-            actions->execute(block_to_fill_as_result);
-        }
+        if (default_actions)
+            default_actions->execute(block_to_fill_as_result);
 
         if (block_to_fill_as_result.has(ProtonConsts::RESERVED_COLUMN_NAMES[0])
             && block_to_fill_as_result.has(ProtonConsts::RESERVED_COLUMN_NAMES[1]))
@@ -528,7 +552,6 @@ class GenerateRandomSource final : public ISource
     Block block_full;
     Block block_to_fill;
     const ColumnsDescription our_columns;
-    pcg64 rng;
     ContextPtr context;
     Int64 boundary_time;
     UInt64 block_idx_in_window;
@@ -546,6 +569,9 @@ class GenerateRandomSource final : public ISource
     UInt64 total_events;
     UInt64 generated_events = 0;
     Poco::Logger * log;
+    std::shared_ptr<ExpressionActions> default_actions = nullptr;
+    // <shard_num, sequence_num, rng>
+    std::tuple<Int64, Int32, pcg64> data_generate_helper;
 
     static Block & prepareBlockToFill(Block & block)
     {
@@ -717,7 +743,8 @@ Pipe StorageRandom::read(
                     0,
                     1000,
                     query_info.syntax_analyzer_result->streaming,
-                    events_share));
+                    events_share,
+                    i));
             }
 
             pipes.emplace_back(std::make_shared<GenerateRandomSource>(
@@ -729,7 +756,8 @@ Pipe StorageRandom::read(
                 0,
                 1000,
                 query_info.syntax_analyzer_result->streaming,
-                events_share + events_remainder));
+                events_share + events_remainder,
+                shards - 1));
         }
         else
         {
@@ -743,7 +771,8 @@ Pipe StorageRandom::read(
                 eps,
                 1000,
                 query_info.syntax_analyzer_result->streaming,
-                max_events));
+                max_events,
+                0));
         }
     }
     else
@@ -762,7 +791,8 @@ Pipe StorageRandom::read(
                 eps_thread,
                 interval_time,
                 query_info.syntax_analyzer_result->streaming,
-                events_share));
+                events_share,
+                i));
         }
 
         /// The last thread will do the remaining work
@@ -775,9 +805,17 @@ Pipe StorageRandom::read(
             eps_thread + remainder,
             interval_time,
             query_info.syntax_analyzer_result->streaming,
-            events_share + events_remainder));
+            events_share + events_remainder,
+            shards - 1));
     }
     return Pipe::unitePipes(std::move(pipes));
 }
 
+NamesAndTypesList StorageRandom::getVirtuals() const
+{
+    return NamesAndTypesList {
+        {ProtonConsts::RESERVED_EVENT_SEQUENCE_ID, std::make_shared<DataTypeInt64>()},
+        {ProtonConsts::RESERVED_SHARD, std::make_shared<DataTypeInt32>()},
+    };
+}
 }
diff --git a/src/Storages/Streaming/StorageRandom.h b/src/Storages/Streaming/StorageRandom.h
@@ -59,6 +59,8 @@ class StorageRandom final : public shared_ptr_helper<StorageRandom>, public ISto
         size_t max_block_size,
         size_t num_streams) override;
 
+    NamesAndTypesList getVirtuals() const override;
+
     bool supportsStreamingQuery() const override { return true; }
     bool hasEvenlyDistributedRead() const override { return true; }
 
diff --git a/tests/stream/test_stream_smoke/0021_random_stream.json b/tests/stream/test_stream_smoke/0021_random_stream.json
