From: Kaan Date: Wed, 19 Mar 2025 15:58:30 +0000 (-0400) Subject: Implement GROUPS frame spec for window functions X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=7e28adbe0c965645affe23e57cf99aa6e16a24e5;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git Implement GROUPS frame spec for window functions Implemented support for the GROUPS frame specification in window functions by adding :paramref:`_sql.over.groups` option to :func:`_sql.over` and :meth:`.FunctionElement.over`. Pull request courtesy Kaan Dikmen. Fixes: #12450 Closes: #12445 Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/12445 Pull-request-sha: c0808e135f15c7fef3a3abcf28465673f38eb428 Change-Id: I9ff504a9c9650485830c4a0eaf44162898a3a2ad --- diff --git a/doc/build/changelog/unreleased_20/12450.rst b/doc/build/changelog/unreleased_20/12450.rst new file mode 100644 index 0000000000..dde46985a5 --- /dev/null +++ b/doc/build/changelog/unreleased_20/12450.rst @@ -0,0 +1,7 @@ +.. change:: + :tags: sql, usecase + :tickets: 12450 + + Implemented support for the GROUPS frame specification in window functions + by adding :paramref:`_sql.over.groups` option to :func:`_sql.over` + and :meth:`.FunctionElement.over`. Pull request courtesy Kaan Dikmen. diff --git a/lib/sqlalchemy/sql/_elements_constructors.py b/lib/sqlalchemy/sql/_elements_constructors.py index 799c87c82b..b5f3c74515 100644 --- a/lib/sqlalchemy/sql/_elements_constructors.py +++ b/lib/sqlalchemy/sql/_elements_constructors.py @@ -1500,6 +1500,7 @@ def over( order_by: Optional[_ByArgument] = None, range_: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, rows: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, + groups: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, ) -> Over[_T]: r"""Produce an :class:`.Over` object against a function. @@ -1517,8 +1518,9 @@ def over( ROW_NUMBER() OVER(ORDER BY some_column) - Ranges are also possible using the :paramref:`.expression.over.range_` - and :paramref:`.expression.over.rows` parameters. These + Ranges are also possible using the :paramref:`.expression.over.range_`, + :paramref:`.expression.over.rows`, and :paramref:`.expression.over.groups` + parameters. These mutually-exclusive parameters each accept a 2-tuple, which contains a combination of integers and None:: @@ -1551,6 +1553,10 @@ def over( func.row_number().over(order_by="x", range_=(1, 3)) + * GROUPS BETWEEN 1 FOLLOWING AND 3 FOLLOWING:: + + func.row_number().over(order_by="x", groups=(1, 3)) + :param element: a :class:`.FunctionElement`, :class:`.WithinGroup`, or other compatible construct. :param partition_by: a column element or string, or a list @@ -1562,10 +1568,14 @@ def over( :param range\_: optional range clause for the window. This is a tuple value which can contain integer values or ``None``, and will render a RANGE BETWEEN PRECEDING / FOLLOWING clause. - :param rows: optional rows clause for the window. This is a tuple value which can contain integer values or None, and will render a ROWS BETWEEN PRECEDING / FOLLOWING clause. + :param groups: optional groups clause for the window. This is a + tuple value which can contain integer values or ``None``, + and will render a GROUPS BETWEEN PRECEDING / FOLLOWING clause. + + .. versionadded:: 2.0.40 This function is also available from the :data:`~.expression.func` construct itself via the :meth:`.FunctionElement.over` method. @@ -1579,7 +1589,7 @@ def over( :func:`_expression.within_group` """ # noqa: E501 - return Over(element, partition_by, order_by, range_, rows) + return Over(element, partition_by, order_by, range_, rows, groups) @_document_text_coercion("text", ":func:`.text`", ":paramref:`.text.text`") diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py index 79dd71ccf9..cdcf9f5c72 100644 --- a/lib/sqlalchemy/sql/compiler.py +++ b/lib/sqlalchemy/sql/compiler.py @@ -2880,6 +2880,8 @@ class SQLCompiler(Compiled): range_ = f"RANGE BETWEEN {self.process(over.range_, **kwargs)}" elif over.rows is not None: range_ = f"ROWS BETWEEN {self.process(over.rows, **kwargs)}" + elif over.groups is not None: + range_ = f"GROUPS BETWEEN {self.process(over.groups, **kwargs)}" else: range_ = None diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py index c9aac427db..42dfe61106 100644 --- a/lib/sqlalchemy/sql/elements.py +++ b/lib/sqlalchemy/sql/elements.py @@ -4212,6 +4212,7 @@ class Over(ColumnElement[_T]): ("partition_by", InternalTraversal.dp_clauseelement), ("range_", InternalTraversal.dp_clauseelement), ("rows", InternalTraversal.dp_clauseelement), + ("groups", InternalTraversal.dp_clauseelement), ] order_by: Optional[ClauseList] = None @@ -4223,6 +4224,7 @@ class Over(ColumnElement[_T]): range_: Optional[_FrameClause] rows: Optional[_FrameClause] + groups: Optional[_FrameClause] def __init__( self, @@ -4231,6 +4233,7 @@ class Over(ColumnElement[_T]): order_by: Optional[_ByArgument] = None, range_: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, rows: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, + groups: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, ): self.element = element if order_by is not None: @@ -4243,19 +4246,14 @@ class Over(ColumnElement[_T]): _literal_as_text_role=roles.ByOfRole, ) - if range_: - self.range_ = _FrameClause(range_) - if rows: - raise exc.ArgumentError( - "'range_' and 'rows' are mutually exclusive" - ) - else: - self.rows = None - elif rows: - self.rows = _FrameClause(rows) - self.range_ = None + if sum(bool(item) for item in (range_, rows, groups)) > 1: + raise exc.ArgumentError( + "only one of 'rows', 'range_', or 'groups' may be provided" + ) else: - self.rows = self.range_ = None + self.range_ = _FrameClause(range_) if range_ else None + self.rows = _FrameClause(rows) if rows else None + self.groups = _FrameClause(groups) if groups else None if not TYPE_CHECKING: @@ -4409,6 +4407,7 @@ class WithinGroup(ColumnElement[_T]): order_by: Optional[_ByArgument] = None, rows: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, range_: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, + groups: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, ) -> Over[_T]: """Produce an OVER clause against this :class:`.WithinGroup` construct. @@ -4423,6 +4422,7 @@ class WithinGroup(ColumnElement[_T]): order_by=order_by, range_=range_, rows=rows, + groups=groups, ) @overload @@ -4540,6 +4540,7 @@ class FunctionFilter(Generative, ColumnElement[_T]): ] = None, range_: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, rows: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, + groups: Optional[typing_Tuple[Optional[int], Optional[int]]] = None, ) -> Over[_T]: """Produce an OVER clause against this filtered function. @@ -4565,6 +4566,7 @@ class FunctionFilter(Generative, ColumnElement[_T]): order_by=order_by, range_=range_, rows=rows, + groups=groups, ) def within_group( diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py index 87a68cfd90..7148d28281 100644 --- a/lib/sqlalchemy/sql/functions.py +++ b/lib/sqlalchemy/sql/functions.py @@ -435,6 +435,7 @@ class FunctionElement(Executable, ColumnElement[_T], FromClause, Generative): order_by: Optional[_ByArgument] = None, rows: Optional[Tuple[Optional[int], Optional[int]]] = None, range_: Optional[Tuple[Optional[int], Optional[int]]] = None, + groups: Optional[Tuple[Optional[int], Optional[int]]] = None, ) -> Over[_T]: """Produce an OVER clause against this function. @@ -466,6 +467,7 @@ class FunctionElement(Executable, ColumnElement[_T], FromClause, Generative): order_by=order_by, rows=rows, range_=range_, + groups=groups, ) def within_group( diff --git a/test/ext/test_serializer.py b/test/ext/test_serializer.py index 40544f3ba0..fb92c752a6 100644 --- a/test/ext/test_serializer.py +++ b/test/ext/test_serializer.py @@ -301,6 +301,16 @@ class SerializeTest(AssertsCompiledSQL, fixtures.MappedTest): "max(users.name) OVER (ROWS BETWEEN CURRENT " "ROW AND UNBOUNDED FOLLOWING)", ), + ( + lambda: func.max(users.c.name).over(groups=(None, 0)), + "max(users.name) OVER (GROUPS BETWEEN UNBOUNDED " + "PRECEDING AND CURRENT ROW)", + ), + ( + lambda: func.max(users.c.name).over(groups=(0, None)), + "max(users.name) OVER (GROUPS BETWEEN CURRENT " + "ROW AND UNBOUNDED FOLLOWING)", + ), ) def test_over(self, over_fn, sql): o = over_fn() diff --git a/test/sql/test_compare.py b/test/sql/test_compare.py index c42bdac7c1..733dcd0aeb 100644 --- a/test/sql/test_compare.py +++ b/test/sql/test_compare.py @@ -452,6 +452,7 @@ class CoreFixtures: func.row_number().over(order_by=table_a.c.a, range_=(0, 10)), func.row_number().over(order_by=table_a.c.a, range_=(None, 10)), func.row_number().over(order_by=table_a.c.a, rows=(None, 20)), + func.row_number().over(order_by=table_a.c.a, groups=(None, 20)), func.row_number().over(order_by=table_a.c.b), func.row_number().over( order_by=table_a.c.a, partition_by=table_a.c.b @@ -1202,6 +1203,14 @@ class CoreFixtures: order_by=table_a.c.a, range_=(random.randint(50, 60), None), ), + func.row_number().over( + order_by=table_a.c.a, + groups=(random.randint(50, 60), random.randint(60, 70)), + ), + func.row_number().over( + order_by=table_a.c.a, + groups=(random.randint(-40, -20), random.randint(60, 70)), + ), ) dont_compare_values_fixtures.append(_numeric_agnostic_window_functions) diff --git a/test/sql/test_compiler.py b/test/sql/test_compiler.py index 5995c5848f..5e86e14db7 100644 --- a/test/sql/test_compiler.py +++ b/test/sql/test_compiler.py @@ -3209,6 +3209,41 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL): checkparams={"param_1": 10, "param_2": 1}, ) + self.assert_compile( + select(func.row_number().over(order_by=expr, groups=(None, 0))), + "SELECT row_number() OVER " + "(ORDER BY mytable.myid GROUPS BETWEEN " + "UNBOUNDED PRECEDING AND CURRENT ROW)" + " AS anon_1 FROM mytable", + ) + + self.assert_compile( + select(func.row_number().over(order_by=expr, groups=(-5, 10))), + "SELECT row_number() OVER " + "(ORDER BY mytable.myid GROUPS BETWEEN " + ":param_1 PRECEDING AND :param_2 FOLLOWING)" + " AS anon_1 FROM mytable", + checkparams={"param_1": 5, "param_2": 10}, + ) + + self.assert_compile( + select(func.row_number().over(order_by=expr, groups=(1, 10))), + "SELECT row_number() OVER " + "(ORDER BY mytable.myid GROUPS BETWEEN " + ":param_1 FOLLOWING AND :param_2 FOLLOWING)" + " AS anon_1 FROM mytable", + checkparams={"param_1": 1, "param_2": 10}, + ) + + self.assert_compile( + select(func.row_number().over(order_by=expr, groups=(-10, -1))), + "SELECT row_number() OVER " + "(ORDER BY mytable.myid GROUPS BETWEEN " + ":param_1 PRECEDING AND :param_2 PRECEDING)" + " AS anon_1 FROM mytable", + checkparams={"param_1": 10, "param_2": 1}, + ) + def test_over_invalid_framespecs(self): assert_raises_message( exc.ArgumentError, @@ -3226,10 +3261,35 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL): assert_raises_message( exc.ArgumentError, - "'range_' and 'rows' are mutually exclusive", + "only one of 'rows', 'range_', or 'groups' may be provided", + func.row_number().over, + range_=(-5, 8), + rows=(-2, 5), + ) + + assert_raises_message( + exc.ArgumentError, + "only one of 'rows', 'range_', or 'groups' may be provided", + func.row_number().over, + range_=(-5, 8), + groups=(None, None), + ) + + assert_raises_message( + exc.ArgumentError, + "only one of 'rows', 'range_', or 'groups' may be provided", + func.row_number().over, + rows=(-2, 5), + groups=(None, None), + ) + + assert_raises_message( + exc.ArgumentError, + "only one of 'rows', 'range_', or 'groups' may be provided", func.row_number().over, range_=(-5, 8), rows=(-2, 5), + groups=(None, None), ) def test_over_within_group(self): diff --git a/test/sql/test_functions.py b/test/sql/test_functions.py index 163df0a0d7..28cdb03a96 100644 --- a/test/sql/test_functions.py +++ b/test/sql/test_functions.py @@ -844,6 +844,34 @@ class CompileTest(fixtures.TestBase, AssertsCompiledSQL): "AS anon_1 FROM mytable", ) + def test_funcfilter_windowing_groups(self): + self.assert_compile( + select( + func.rank() + .filter(table1.c.name > "foo") + .over(groups=(1, 5), partition_by=["description"]) + ), + "SELECT rank() FILTER (WHERE mytable.name > :name_1) " + "OVER (PARTITION BY mytable.description GROUPS BETWEEN :param_1 " + "FOLLOWING AND :param_2 FOLLOWING) " + "AS anon_1 FROM mytable", + ) + + def test_funcfilter_windowing_groups_positional(self): + self.assert_compile( + select( + func.rank() + .filter(table1.c.name > "foo") + .over(groups=(1, 5), partition_by=["description"]) + ), + "SELECT rank() FILTER (WHERE mytable.name > ?) " + "OVER (PARTITION BY mytable.description GROUPS BETWEEN ? " + "FOLLOWING AND ? FOLLOWING) " + "AS anon_1 FROM mytable", + checkpositional=("foo", 1, 5), + dialect="default_qmark", + ) + def test_funcfilter_more_criteria(self): ff = func.rank().filter(table1.c.name > "foo") ff2 = ff.filter(table1.c.myid == 1)