ISSUES-2553 add offset parameter to numbers

dbms/src/Storages/System/StorageSystemNumbers.cpp

@@ -5,7 +5,6 @@
 #include <DataStreams/LimitBlockInputStream.h>
 #include <Storages/System/StorageSystemNumbers.h>
 
-
 namespace DB
 {
 
@@ -44,8 +43,8 @@ private:
 };
 
 
-StorageSystemNumbers::StorageSystemNumbers(const std::string & name_, bool multithreaded_, size_t limit_)
-    : name(name_), multithreaded(multithreaded_), limit(limit_)
+StorageSystemNumbers::StorageSystemNumbers(const std::string & name_, bool multithreaded_, size_t limit_, size_t offset_)
+    : name(name_), multithreaded(multithreaded_), limit(limit_), offset(offset_)
 {
     setColumns(ColumnsDescription({{"number", std::make_shared<DataTypeUInt64>()}}));
 }
@@ -74,7 +73,7 @@ BlockInputStreams StorageSystemNumbers::read(
     BlockInputStreams res(num_streams);
     for (size_t i = 0; i < num_streams; ++i)
     {
-        res[i] = std::make_shared<NumbersBlockInputStream>(max_block_size, i * max_block_size, num_streams * max_block_size);
+        res[i] = std::make_shared<NumbersBlockInputStream>(max_block_size, offset + i * max_block_size, num_streams * max_block_size);
 
         if (limit)  /// This formula is how to split 'limit' elements to 'num_streams' chunks almost uniformly.
             res[i] = std::make_shared<LimitBlockInputStream>(res[i], limit * (i + 1) / num_streams - limit * i / num_streams, 0);

dbms/src/Storages/System/StorageSystemNumbers.h

@@ -37,10 +37,11 @@ private:
     const std::string name;
     bool multithreaded;
     size_t limit;
+    size_t offset;
 
 protected:
     /// limit: 0 means unlimited.
-    StorageSystemNumbers(const std::string & name_, bool multithreaded_, size_t limit_ = 0);
+    StorageSystemNumbers(const std::string & name_, bool multithreaded_, size_t limit_ = 0, size_t offset_ = 0);
 };
 
 }

dbms/src/TableFunctions/TableFunctionNumbers.cpp

@@ -3,7 +3,6 @@
 #include <TableFunctions/TableFunctionFactory.h>
 #include <Parsers/ASTFunction.h>
 #include <Parsers/ASTLiteral.h>
-#include <Common/Exception.h>
 #include <Common/typeid_cast.h>
 #include <Storages/System/StorageSystemNumbers.h>
 #include <Interpreters/evaluateConstantExpression.h>
@@ -20,31 +19,33 @@ namespace ErrorCodes
 
 StoragePtr TableFunctionNumbers::executeImpl(const ASTPtr & ast_function, const Context & context) const
 {
-    ASTs & args_func = typeid_cast<ASTFunction &>(*ast_function).children;
+    if (const ASTFunction * function = typeid_cast<ASTFunction *>(ast_function.get()))
+    {
+        auto arguments = function->arguments->children;
 
-    if (args_func.size() != 1)
-        throw Exception("Table function 'numbers' requires exactly one argument: amount of numbers.",
-            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+        if (arguments.size() != 1 && arguments.size() != 2)
+            throw Exception("Table function 'numbers' requires 'length' or 'offset, length'.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
 
-    ASTs & args = typeid_cast<ASTExpressionList &>(*args_func.at(0)).children;
 
-    if (args.size() != 1)
-        throw Exception("Table function 'numbers' requires exactly one argument: amount of numbers.",
-            ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
+        UInt64 offset = arguments.size() == 2 ? evaluateArgument(context, arguments[0]) : 0;
+        UInt64 length = arguments.size() == 2 ? evaluateArgument(context, arguments[1]) : evaluateArgument(context, arguments[0]);
 
-    args[0] = evaluateConstantExpressionOrIdentifierAsLiteral(args[0], context);
-
-    UInt64 limit = static_cast<const ASTLiteral &>(*args[0]).value.safeGet<UInt64>();
-
-    auto res = StorageSystemNumbers::create(getName(), false, limit);
-    res->startup();
-    return res;
+        auto res = StorageSystemNumbers::create(getName(), false, length, offset);
+        res->startup();
+        return res;
+    }
+    throw new Exception("Table function 'numbers' requires 'limit' or 'offset, limit'.", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
 }
 
-
 void registerTableFunctionNumbers(TableFunctionFactory & factory)
 {
     factory.registerFunction<TableFunctionNumbers>();
 }
 
+
+UInt64 TableFunctionNumbers::evaluateArgument(const Context & context, ASTPtr & argument) const
+{
+    return static_cast<const ASTLiteral &>(*evaluateConstantExpressionOrIdentifierAsLiteral(argument, context)).value.safeGet<UInt64>();
+}
+
 }

dbms/src/TableFunctions/TableFunctionNumbers.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <TableFunctions/ITableFunction.h>
+#include <Core/Types.h>
 
 
 namespace DB
@@ -17,6 +18,8 @@ public:
     std::string getName() const override { return name; }
 private:
     StoragePtr executeImpl(const ASTPtr & ast_function, const Context & context) const override;
+
+    UInt64 evaluateArgument(const Context & context, ASTPtr & argument) const;
 };
 
 

dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.reference

@@ -0,0 +1 @@
+1

dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.sql

@@ -0,0 +1,2 @@
+SET max_rows_to_read = 1;
+SELECT * FROM numbers(1, 1);