From: Kaan Date: Wed, 19 Mar 2025 15:58:30 +0000 (-0400) Subject: Implement GROUPS frame spec for window functions X-Git-Tag: rel_2_0_40~4 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=1dedb8bf9bbdbab2317f484b8c26e9ee58480370;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 (cherry picked from commit 0bbc515f904446d3f0beede54321b628f32fbdad) --- 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 b628fcc9b5..3359998f3d 100644 --- a/lib/sqlalchemy/sql/_elements_constructors.py +++ b/lib/sqlalchemy/sql/_elements_constructors.py @@ -1508,6 +1508,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. @@ -1525,8 +1526,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:: @@ -1559,6 +1561,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 @@ -1570,10 +1576,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. @@ -1587,7 +1597,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 9130cdb2c3..098667f92b 100644 --- a/lib/sqlalchemy/sql/compiler.py +++ b/lib/sqlalchemy/sql/compiler.py @@ -2896,14 +2896,18 @@ class SQLCompiler(Compiled): def visit_over(self, over, **kwargs): text = over.element._compiler_dispatch(self, **kwargs) - if over.range_: + if over.range_ is not None: range_ = "RANGE BETWEEN %s" % self._format_frame_clause( over.range_, **kwargs ) - elif over.rows: + elif over.rows is not None: range_ = "ROWS BETWEEN %s" % self._format_frame_clause( over.rows, **kwargs ) + elif over.groups is not None: + range_ = "GROUPS BETWEEN %s" % self._format_frame_clause( + over.groups, **kwargs + ) else: range_ = None diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py index b259f96463..88cb2529d8 100644 --- a/lib/sqlalchemy/sql/elements.py +++ b/lib/sqlalchemy/sql/elements.py @@ -4194,6 +4194,7 @@ class Over(ColumnElement[_T]): ("partition_by", InternalTraversal.dp_clauseelement), ("range_", InternalTraversal.dp_plain_obj), ("rows", InternalTraversal.dp_plain_obj), + ("groups", InternalTraversal.dp_plain_obj), ] order_by: Optional[ClauseList] = None @@ -4205,6 +4206,7 @@ class Over(ColumnElement[_T]): range_: Optional[typing_Tuple[_IntOrRange, _IntOrRange]] rows: Optional[typing_Tuple[_IntOrRange, _IntOrRange]] + groups: Optional[typing_Tuple[_IntOrRange, _IntOrRange]] def __init__( self, @@ -4213,6 +4215,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: @@ -4225,19 +4228,14 @@ class Over(ColumnElement[_T]): _literal_as_text_role=roles.ByOfRole, ) - if range_: - self.range_ = self._interpret_range(range_) - if rows: - raise exc.ArgumentError( - "'range_' and 'rows' are mutually exclusive" - ) - else: - self.rows = None - elif rows: - self.rows = self._interpret_range(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_ = self._interpret_range(range_) if range_ else None + self.rows = self._interpret_range(rows) if rows else None + self.groups = self._interpret_range(groups) if groups else None def __reduce__(self): return self.__class__, ( @@ -4246,6 +4244,7 @@ class Over(ColumnElement[_T]): self.order_by, self.range_, self.rows, + self.groups, ) def _interpret_range( @@ -4360,6 +4359,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. @@ -4374,6 +4374,7 @@ class WithinGroup(ColumnElement[_T]): order_by=order_by, range_=range_, rows=rows, + groups=groups, ) @overload @@ -4491,6 +4492,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. @@ -4516,6 +4518,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 ea02279d48..b0f6655cb1 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 96e3e7661d..0457770442 100644 --- a/test/sql/test_compare.py +++ b/test/sql/test_compare.py @@ -440,6 +440,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 diff --git a/test/sql/test_compiler.py b/test/sql/test_compiler.py index 12ba24e170..7c43e60db3 100644 --- a/test/sql/test_compiler.py +++ b/test/sql/test_compiler.py @@ -3213,6 +3213,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, @@ -3230,10 +3265,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)