feat(duckdb): implement NthValue

ibis-project · May 23, 2022 · 1bf2842 · 1bf2842
1 parent 37ff328
commit 1bf2842
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/ibis/backends/base/sql/alchemy/registry.py b/ibis/backends/base/sql/alchemy/registry.py
@@ -580,6 +580,11 @@ def _string_join(t, expr):
 }
 
 
+def _nth_value(t, expr):
+    op = expr.op()
+    return sa.func.nth_value(t.translate(op.arg), t.translate(op.nth) + 1)
+
+
 sqlalchemy_window_functions_registry = {
     ops.Lag: _lag,
     ops.Lead: _lead,
@@ -590,6 +595,7 @@ def _string_join(t, expr):
     ops.DenseRank: unary(lambda arg: sa.func.dense_rank()),
     ops.MinRank: unary(lambda arg: sa.func.rank()),
     ops.PercentRank: unary(lambda arg: sa.func.percent_rank()),
+    ops.NthValue: _nth_value,
     ops.Window: _window,
     ops.CumulativeOp: _window,
     ops.CumulativeMax: unary(sa.func.max),

diff --git a/ibis/backends/tests/test_window.py b/ibis/backends/tests/test_window.py
@@ -69,6 +69,23 @@ def calc_zscore(s):
             lambda t: t.float_col.transform('last'),
             id='last',
         ),
+        param(
+            lambda t, win: t.float_col.nth(3).over(win),
+            lambda t: t.float_col.apply(
+                lambda s: pd.concat(
+                    [
+                        pd.Series(np.nan, index=s.index[:3], dtype="float32"),
+                        pd.Series(
+                            s.iloc[3],
+                            index=s.index[3:],
+                            dtype="float32",
+                        ),
+                    ]
+                )
+            ),
+            id="nth",
+            marks=pytest.mark.notimpl(["pandas"]),
+        ),
         param(
             lambda _, win: ibis.row_number().over(win),
             lambda t: t.cumcount(),