diff --git a/dbms/src/Storages/System/StorageSystemNumbers.cpp b/dbms/src/Storages/System/StorageSystemNumbers.cpp index 49a3c7ca513..367a2639870 100644 --- a/dbms/src/Storages/System/StorageSystemNumbers.cpp +++ b/dbms/src/Storages/System/StorageSystemNumbers.cpp @@ -5,7 +5,6 @@ #include #include - 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()}})); } @@ -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(max_block_size, i * max_block_size, num_streams * max_block_size); + res[i] = std::make_shared(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(res[i], limit * (i + 1) / num_streams - limit * i / num_streams, 0); diff --git a/dbms/src/Storages/System/StorageSystemNumbers.h b/dbms/src/Storages/System/StorageSystemNumbers.h index e0769fb3968..30c68cbd853 100644 --- a/dbms/src/Storages/System/StorageSystemNumbers.h +++ b/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); }; } diff --git a/dbms/src/TableFunctions/TableFunctionNumbers.cpp b/dbms/src/TableFunctions/TableFunctionNumbers.cpp index 60136dccfab..1970a757b2d 100644 --- a/dbms/src/TableFunctions/TableFunctionNumbers.cpp +++ b/dbms/src/TableFunctions/TableFunctionNumbers.cpp @@ -3,7 +3,6 @@ #include #include #include -#include #include #include #include @@ -20,31 +19,33 @@ namespace ErrorCodes StoragePtr TableFunctionNumbers::executeImpl(const ASTPtr & ast_function, const Context & context) const { - ASTs & args_func = typeid_cast(*ast_function).children; + if (const ASTFunction * function = typeid_cast(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(*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(*args[0]).value.safeGet(); - - 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(); } + +UInt64 TableFunctionNumbers::evaluateArgument(const Context & context, ASTPtr & argument) const +{ + return static_cast(*evaluateConstantExpressionOrIdentifierAsLiteral(argument, context)).value.safeGet(); +} + } diff --git a/dbms/src/TableFunctions/TableFunctionNumbers.h b/dbms/src/TableFunctions/TableFunctionNumbers.h index 6a3ab7f9090..ed060a6450a 100644 --- a/dbms/src/TableFunctions/TableFunctionNumbers.h +++ b/dbms/src/TableFunctions/TableFunctionNumbers.h @@ -1,6 +1,7 @@ #pragma once #include +#include 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; }; diff --git a/dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.reference b/dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.reference new file mode 100644 index 00000000000..d00491fd7e5 --- /dev/null +++ b/dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.reference @@ -0,0 +1 @@ +1 diff --git a/dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.sql b/dbms/tests/queries/0_stateless/00647_select_numbers_with_offset.sql new file mode 100644 index 00000000000..3cea011a45d --- /dev/null +++ b/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);