From 65e63235ec953df55698abec692eeb7560b0e92d Mon Sep 17 00:00:00 2001 From: vdimir Date: Mon, 24 Apr 2023 10:41:43 +0000 Subject: [PATCH] Custom default window frame for ntile function --- src/Processors/Transforms/WindowTransform.cpp | 54 +++++++++++-------- .../0_stateless/02560_window_ntile.reference | 39 ++++++++++---- .../0_stateless/02560_window_ntile.sql | 21 ++++---- 3 files changed, 74 insertions(+), 40 deletions(-) diff --git a/src/Processors/Transforms/WindowTransform.cpp b/src/Processors/Transforms/WindowTransform.cpp index 04a1f12f30a..d2cce2cb890 100644 --- a/src/Processors/Transforms/WindowTransform.cpp +++ b/src/Processors/Transforms/WindowTransform.cpp @@ -44,6 +44,8 @@ public: // Must insert the result for current_row. virtual void windowInsertResultInto(const WindowTransform * transform, size_t function_index) = 0; + + virtual std::optional getDefaultFrame() const { return {}; } }; // Compares ORDER BY column values at given rows to find the boundaries of frame: @@ -222,6 +224,15 @@ WindowTransform::WindowTransform(const Block & input_header_, /// Currently we have slightly wrong mixup of the interfaces of Window and Aggregate functions. workspace.window_function_impl = dynamic_cast(const_cast(aggregate_function.get())); + /// Some functions may have non-standard default frame. + /// Use it if it's the only function over the current window. + if (window_description.frame.is_default && functions.size() == 1 && workspace.window_function_impl) + { + auto custom_default_frame = workspace.window_function_impl->getDefaultFrame(); + if (custom_default_frame) + window_description.frame = *custom_default_frame; + } + workspace.aggregate_function_state.reset( aggregate_function->sizeOfData(), aggregate_function->alignOfData()); @@ -1977,18 +1988,23 @@ struct WindowFunctionNtile final : public WindowFunction : WindowFunction(name_, argument_types_, parameters_, std::make_shared()) { if (argument_types.size() != 1) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "Function {} takes exactly one parameter", name_); - } + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Function {} takes exactly one argument", name_); + auto type_id = argument_types[0]->getTypeId(); if (type_id != TypeIndex::UInt8 && type_id != TypeIndex::UInt16 && type_id != TypeIndex::UInt32 && type_id != TypeIndex::UInt64) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's argument type must be an unsigned integer (not larger then 64-bit), but got {}", argument_types[0]->getName()); - } + throw Exception(ErrorCodes::BAD_ARGUMENTS, "'{}' argument type must be an unsigned integer (not larger than 64-bit), got {}", name_, argument_types[0]->getName()); } bool allocatesMemoryInArena() const override { return false; } + std::optional getDefaultFrame() const override + { + WindowFrame frame; + frame.type = WindowFrame::FrameType::ROWS; + frame.end_type = WindowFrame::BoundaryType::Unbounded; + return frame; + } + void windowInsertResultInto(const WindowTransform * transform, size_t function_index) override { @@ -1999,7 +2015,7 @@ struct WindowFunctionNtile final : public WindowFunction const auto & workspace = transform->workspaces[function_index]; const auto & arg_col = *current_block.original_input_columns[workspace.argument_column_indices[0]]; if (!isColumnConst(arg_col)) - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's argument must be a constant"); + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Argument of 'ntile' function must be a constant"); auto type_id = argument_types[0]->getTypeId(); if (type_id == TypeIndex::UInt8) buckets = arg_col[transform->current_row.row].get(); @@ -2012,7 +2028,7 @@ struct WindowFunctionNtile final : public WindowFunction if (!buckets) { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's argument must > 0"); + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Argument of 'ntile' funtcion must be greater than zero"); } } // new partition @@ -2090,22 +2106,16 @@ private: static void checkWindowFrameType(const WindowTransform * transform) { if (transform->order_by_indices.empty()) - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's window frame must have order by clause"); - if (transform->window_description.frame.type != WindowFrame::FrameType::ROWS) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's frame type must be ROWS"); - } - if (transform->window_description.frame.begin_type != WindowFrame::BoundaryType::Unbounded) - { - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's frame start type must be UNBOUNDED PRECEDING"); - } + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Window frame for 'ntile' function must have ORDER BY clause"); - if (transform->window_description.frame.end_type != WindowFrame::BoundaryType::Unbounded) + // We must wait all for the partition end and get the total rows number in this + // partition. So before the end of this partition, there is no any block could be + // dropped out. + bool is_frame_supported = transform->window_description.frame.begin_type == WindowFrame::BoundaryType::Unbounded + && transform->window_description.frame.end_type == WindowFrame::BoundaryType::Unbounded; + if (!is_frame_supported) { - // We must wait all for the partition end and get the total rows number in this - // partition. So before the end of this partition, there is no any block could be - // dropped out. - throw Exception(ErrorCodes::BAD_ARGUMENTS, "ntile's frame end type must be UNBOUNDED FOLLOWING"); + throw Exception(ErrorCodes::BAD_ARGUMENTS, "Window frame for function 'ntile' should be 'ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'"); } } }; diff --git a/tests/queries/0_stateless/02560_window_ntile.reference b/tests/queries/0_stateless/02560_window_ntile.reference index cae0586fa8c..1045fc1011a 100644 --- a/tests/queries/0_stateless/02560_window_ntile.reference +++ b/tests/queries/0_stateless/02560_window_ntile.reference @@ -22,7 +22,28 @@ select a, b, ntile(3) over (partition by a order by b rows between unbounded pre 1 7 3 1 8 3 1 9 3 -select a, b, ntile(2) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(3) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +0 0 1 +0 1 1 +0 2 1 +0 3 1 +0 4 2 +0 5 2 +0 6 2 +0 7 3 +0 8 3 +0 9 3 +1 0 1 +1 1 1 +1 2 1 +1 3 1 +1 4 2 +1 5 2 +1 6 2 +1 7 3 +1 8 3 +1 9 3 +select a, b, ntile(2) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); 0 0 1 0 1 1 0 2 1 @@ -43,7 +64,7 @@ select a, b, ntile(2) over (partition by a order by b rows between unbounded pre 1 7 2 1 8 2 1 9 2 -select a, b, ntile(1) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(1) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); 0 0 1 0 1 1 0 2 1 @@ -64,7 +85,7 @@ select a, b, ntile(1) over (partition by a order by b rows between unbounded pre 1 7 1 1 8 1 1 9 1 -select a, b, ntile(100) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(100) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); 0 0 1 0 1 2 0 2 3 @@ -85,7 +106,7 @@ select a, b, ntile(100) over (partition by a order by b rows between unbounded p 1 7 8 1 8 9 1 9 10 -select a, b, ntile(65535) over (partition by a order by b rows between unbounded preceding and unbounded following) from (select 1 as a, number as b from numbers(65535)) limit 100; +select a, b, ntile(65535) over (partition by a order by b) from (select 1 as a, number as b from numbers(65535)) limit 100; 1 0 1 1 1 2 1 2 3 @@ -187,11 +208,11 @@ select a, b, ntile(65535) over (partition by a order by b rows between unbounded 1 98 99 1 99 100 -- Bad arguments -select a, b, ntile(3.0) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile('2') over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(0) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(-2) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(b + 1) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(3.0) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile('2') over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(0) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(-2) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(b + 1) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -- Bad window type select a, b, ntile(2) over (partition by a) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } select a, b, ntile(2) over (partition by a order by b rows between 4 preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } diff --git a/tests/queries/0_stateless/02560_window_ntile.sql b/tests/queries/0_stateless/02560_window_ntile.sql index 4c25ecf4dd2..f2acf8fc94e 100644 --- a/tests/queries/0_stateless/02560_window_ntile.sql +++ b/tests/queries/0_stateless/02560_window_ntile.sql @@ -2,17 +2,20 @@ -- Normal cases select a, b, ntile(3) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -select a, b, ntile(2) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -select a, b, ntile(1) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -select a, b, ntile(100) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -select a, b, ntile(65535) over (partition by a order by b rows between unbounded preceding and unbounded following) from (select 1 as a, number as b from numbers(65535)) limit 100; +select a, b, ntile(3) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(2) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(1) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(100) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); +select a, b, ntile(65535) over (partition by a order by b) from (select 1 as a, number as b from numbers(65535)) limit 100; + + -- Bad arguments -select a, b, ntile(3.0) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile('2') over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(0) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(-2) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -select a, b, ntile(b + 1) over (partition by a order by b rows between unbounded preceding and unbounded following) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(3.0) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile('2') over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(0) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(-2) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } +select a, b, ntile(b + 1) over (partition by a order by b) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 } -- Bad window type select a, b, ntile(2) over (partition by a) from(select intDiv(number,10) as a, number%10 as b from numbers(20)); -- { serverError 36 }