Add frame exclusion support for window functions

Add the `exclude` parameter to the `Over` construct and all `.over()`
methods, enabling SQL standard frame exclusion clauses
`EXCLUDE CURRENT ROW`, `EXCLUDE GROUP`, `EXCLUDE TIES`,
`EXCLUDE NO OTHERS` in window functions.
Pull request courtesy of Varun Chawla.

Fixes #11671
Closes: #13117
Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/13117
Pull-request-sha: 23f9d34dd80c45dff68ecc8c08331acf22b82279

Change-Id: I8efdb06876d5a11a9f5ed9abec2c187c6c9b7e5e

diff --git a/doc/build/changelog/unreleased_21/11671.rst b/doc/build/changelog/unreleased_21/11671.rst
new file mode 100644 (file)
index 0000000..a9ab9ac
--- /dev/null
+++ b/doc/build/changelog/unreleased_21/11671.rst
@@ -0,0 +1,9 @@
+.. change::
+    :tags: usecase, sql
+    :tickets: 11671
+
+    Add the `exclude` parameter to the `Over` construct and all `.over()`
+    methods, enabling SQL standard frame exclusion clauses
+    `EXCLUDE CURRENT ROW`, `EXCLUDE GROUP`, `EXCLUDE TIES`,
+    `EXCLUDE NO OTHERS` in window functions.
+    Pull request courtesy of Varun Chawla.
\ No newline at end of file

diff --git a/lib/sqlalchemy/sql/_elements_constructors.py b/lib/sqlalchemy/sql/_elements_constructors.py
index 5b23f0fb06e1af53e158c301868d51d86a485a94..756789697962791f049376a080f270a9838d8841 100644 (file)
--- a/lib/sqlalchemy/sql/_elements_constructors.py
+++ b/lib/sqlalchemy/sql/_elements_constructors.py
@@ -1621,6 +1621,7 @@ def over(
     range_: _FrameIntTuple | FrameClause | None = None,
     rows: _FrameIntTuple | FrameClause | None = None,
     groups: _FrameIntTuple | FrameClause | None = None,
+    exclude: str | None = None,
 ) -> Over[_T]:
     r"""Produce an :class:`.Over` object against a function.
 
@@ -1725,6 +1726,14 @@ def over(
 
      .. versionadded:: 2.0.40
 
+    :param exclude: optional string for the frame exclusion clause.
+     This is a string value which can be one of ``CURRENT ROW``,
+     ``GROUP``, ``TIES``, or ``NO OTHERS`` and will render an
+     EXCLUDE clause within the window frame specification.  Requires
+     that one of ``rows``, ``range_``, or ``groups`` is also specified.
+
+     .. versionadded:: 2.1
+
     This function is also available from the :data:`~.expression.func`
     construct itself via the :meth:`.FunctionElement.over` method.
 
@@ -1737,7 +1746,15 @@ def over(
         :func:`_expression.within_group`
 
     """  # noqa: E501
-    return Over(element, partition_by, order_by, range_, rows, groups)
+    return Over(
+        element,
+        partition_by,
+        order_by,
+        range_,
+        rows,
+        groups,
+        exclude,
+    )
 
 
 @_document_text_coercion("text", ":func:`.text`", ":paramref:`.text.text`")

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index f0e608e748da9d6b0d10e85e49fb02538b23be4f..66de4813eef0212e397b0486e6df2608c780d4bd 100644 (file)
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -262,6 +262,9 @@ FK_ON_UPDATE = re.compile(
     r"^(?:RESTRICT|CASCADE|SET NULL|NO ACTION|SET DEFAULT)$", re.I
 )
 FK_INITIALLY = re.compile(r"^(?:DEFERRED|IMMEDIATE)$", re.I)
+_WINDOW_EXCLUDE_RE = re.compile(
+    r"^(?:CURRENT ROW|GROUP|TIES|NO OTHERS)$", re.I
+)
 BIND_PARAMS = re.compile(r"(?<![:\w\$\x5c]):([\w\$]+)(?![:\w\$])", re.UNICODE)
 BIND_PARAMS_ESC = re.compile(r"\x5c(:[\w\$]*)(?![:\w\$])", re.UNICODE)
 
@@ -3000,6 +3003,11 @@ class SQLCompiler(Compiled):
         else:
             range_ = None
 
+        if range_ is not None and over.exclude is not None:
+            range_ += " EXCLUDE " + self.preparer.validate_sql_phrase(
+                over.exclude, _WINDOW_EXCLUDE_RE
+            )
+
         return "%s OVER (%s)" % (
             text,
             " ".join(

diff --git a/lib/sqlalchemy/sql/elements.py b/lib/sqlalchemy/sql/elements.py
index 7343874fff050f4dc6b9daf8e393515e92e971e3..12b9eeefb7fa6527e1c9a7f3fe75a43814ffae91 100644 (file)
--- a/lib/sqlalchemy/sql/elements.py
+++ b/lib/sqlalchemy/sql/elements.py
@@ -4528,6 +4528,7 @@ class Over(ColumnElement[_T]):
         ("range_", InternalTraversal.dp_clauseelement),
         ("rows", InternalTraversal.dp_clauseelement),
         ("groups", InternalTraversal.dp_clauseelement),
+        ("exclude", InternalTraversal.dp_string),
     ]
 
     order_by: Optional[ClauseList] = None
@@ -4540,6 +4541,7 @@ class Over(ColumnElement[_T]):
     range_: FrameClause | None
     rows: FrameClause | None
     groups: FrameClause | None
+    exclude: str | None
 
     def __init__(
         self,
@@ -4549,6 +4551,7 @@ class Over(ColumnElement[_T]):
         range_: _FrameIntTuple | FrameClause | None = None,
         rows: _FrameIntTuple | FrameClause | None = None,
         groups: _FrameIntTuple | FrameClause | None = None,
+        exclude: str | None = None,
     ):
         self.element = element
         if order_by is not None:
@@ -4570,6 +4573,16 @@ class Over(ColumnElement[_T]):
             self.rows = FrameClause._parse(rows, coerce_int=True)
             self.groups = FrameClause._parse(groups, coerce_int=True)
 
+        self.exclude = exclude
+
+        if exclude is not None and (
+            range_ is None and rows is None and groups is None
+        ):
+            raise exc.ArgumentError(
+                "'exclude' requires that one of 'rows', "
+                "'range_', or 'groups' is also specified"
+            )
+
     if not TYPE_CHECKING:
 
         @util.memoized_property
@@ -4804,6 +4817,7 @@ class AggregateOrderBy(WrapsColumnExpression[_T]):
         rows: _FrameIntTuple | FrameClause | None = None,
         range_: _FrameIntTuple | FrameClause | None = None,
         groups: _FrameIntTuple | FrameClause | None = None,
+        exclude: str | None = None,
     ) -> Over[_T]:
         """Produce an OVER clause against this :class:`.WithinGroup`
         construct.
@@ -4819,6 +4833,7 @@ class AggregateOrderBy(WrapsColumnExpression[_T]):
             range_=range_,
             rows=rows,
             groups=groups,
+            exclude=exclude,
         )
 
     @overload
@@ -4948,6 +4963,7 @@ class FunctionFilter(Generative, ColumnElement[_T]):
         range_: _FrameIntTuple | FrameClause | None = None,
         rows: _FrameIntTuple | FrameClause | None = None,
         groups: _FrameIntTuple | FrameClause | None = None,
+        exclude: str | None = None,
     ) -> Over[_T]:
         """Produce an OVER clause against this filtered function.
 
@@ -4974,6 +4990,7 @@ class FunctionFilter(Generative, ColumnElement[_T]):
             range_=range_,
             rows=rows,
             groups=groups,
+            exclude=exclude,
         )
 
     def within_group(

diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
index 58451f320f8c48aa4efe4908d36cbba5ee4f1608..768c8c6856a581927ebcca75137a90da3d15c041 100644 (file)
--- a/lib/sqlalchemy/sql/functions.py
+++ b/lib/sqlalchemy/sql/functions.py
@@ -460,6 +460,7 @@ class FunctionElement(
         rows: _FrameIntTuple | FrameClause | None = None,
         range_: _FrameIntTuple | FrameClause | None = None,
         groups: _FrameIntTuple | FrameClause | None = None,
+        exclude: str | None = None,
     ) -> Over[_T]:
         """Produce an OVER clause against this function.
 
@@ -492,6 +493,7 @@ class FunctionElement(
             rows=rows,
             range_=range_,
             groups=groups,
+            exclude=exclude,
         )
 
     def aggregate_order_by(

diff --git a/test/sql/test_compiler.py b/test/sql/test_compiler.py
index 6d4596e46cea818a551a824a86f085e91a1e4688..cf2afe4448bc1c3d56190d40dfce693fa2e2423d 100644 (file)
--- a/test/sql/test_compiler.py
+++ b/test/sql/test_compiler.py
@@ -3300,6 +3300,98 @@ class SelectTest(fixtures.TestBase, AssertsCompiledSQL):
             checkparams={"param_1": 1, "param_2": 3},
         )
 
+    @testing.combinations(
+        (
+            "rows_current_row",
+            {"rows": (None, 0)},
+            "CURRENT ROW",
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN UNBOUNDED "
+            "PRECEDING AND CURRENT ROW EXCLUDE CURRENT ROW)"
+            " AS anon_1 FROM mytable",
+            {},
+        ),
+        (
+            "range_group",
+            {"range_": (None, 0)},
+            "GROUP",
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid RANGE BETWEEN UNBOUNDED "
+            "PRECEDING AND CURRENT ROW EXCLUDE GROUP)"
+            " AS anon_1 FROM mytable",
+            {},
+        ),
+        (
+            "groups_ties",
+            {"groups": (-1, 1)},
+            "TIES",
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid GROUPS BETWEEN "
+            ":param_1 PRECEDING AND :param_2 FOLLOWING EXCLUDE TIES)"
+            " AS anon_1 FROM mytable",
+            {"param_1": 1, "param_2": 1},
+        ),
+        (
+            "rows_no_others",
+            {"rows": (None, None)},
+            "NO OTHERS",
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN UNBOUNDED "
+            "PRECEDING AND UNBOUNDED FOLLOWING EXCLUDE NO OTHERS)"
+            " AS anon_1 FROM mytable",
+            {},
+        ),
+        (
+            "rows_lowercase_ties",
+            {"rows": (None, 0)},
+            "ties",
+            "SELECT row_number() OVER "
+            "(ORDER BY mytable.myid ROWS BETWEEN UNBOUNDED "
+            "PRECEDING AND CURRENT ROW EXCLUDE ties)"
+            " AS anon_1 FROM mytable",
+            {},
+        ),
+        id_="iaaaa",
+    )
+    def test_over_frame_exclude(
+        self, frame_kwargs, exclude, expected_sql, checkparams
+    ):
+        self.assert_compile(
+            select(
+                func.row_number().over(
+                    order_by=table1.c.myid, exclude=exclude, **frame_kwargs
+                )
+            ),
+            expected_sql,
+            checkparams=checkparams,
+        )
+
+    def test_over_frame_exclude_invalid(self):
+        # invalid exclude value raises at compile time
+        assert_raises_message(
+            exc.CompileError,
+            "Unexpected SQL phrase: 'INVALID'",
+            select(
+                func.row_number().over(
+                    order_by=table1.c.myid,
+                    rows=(None, 0),
+                    exclude="INVALID",
+                )
+            ).compile,
+        )
+
+    def test_over_frame_exclude_requires_frame_spec(self):
+        # exclude without rows/range_/groups raises at construction time
+        with expect_raises_message(
+            exc.ArgumentError,
+            "'exclude' requires that one of 'rows', "
+            "'range_', or 'groups' is also specified",
+        ):
+            func.row_number().over(
+                order_by=table1.c.myid,
+                exclude="CURRENT ROW",
+            )
+
     def test_over_invalid_framespecs(self):
         with expect_raises_message(
             exc.ArgumentError,