Skip to content

Commit 0acc1ad

Browse files
sfc-gh-dryazanovsfc-gh-jdu
sfc-gh-dryazanov
and
sfc-gh-jdu
authored
SNOW-1896935 Implement nth_value (#2936)
Co-authored-by: Jianzhun Du <68252326+sfc-gh-jdu@users.noreply.github.com>
1 parent 8f0f191 commit 0acc1ad

File tree

5 files changed

+264
-1
lines changed

5 files changed

+264
-1
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
- `localtimestamp`
2929
- `max_by`
3030
- `min_by`
31+
- `nth_value`
3132
- `nvl`
3233
- `octet_length`
3334
- `position`

‎docs/source/snowpark/functions.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@ Functions
251251
negate
252252
next_day
253253
not_
254+
nth_value
254255
ntile
255256
nullifzero
256257
nvl

‎src/snowflake/snowpark/_internal/analyzer/window_expression.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -231,3 +231,7 @@ class LastValue(RankRelatedFunctionExpression):
231231

232232
class FirstValue(RankRelatedFunctionExpression):
233233
sql = "FIRST_VALUE"
234+
235+
236+
class NthValue(RankRelatedFunctionExpression):
237+
sql = "NTH_VALUE"

‎src/snowflake/snowpark/functions.py‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,7 @@
184184
Lag,
185185
LastValue,
186186
Lead,
187+
NthValue,
187188
)
188189
from snowflake.snowpark._internal.ast.utils import (
189190
build_builtin_fn_apply,
@@ -8924,13 +8925,41 @@ def first_value(
89248925
ast = proto.Expr()
89258926
build_builtin_fn_apply(ast, "first_value", e, ignore_nulls)
89268927

8927-
c = _to_col_if_str(e, "last_value")
8928+
c = _to_col_if_str(e, "first_value")
89288929

89298930
ans = Column(FirstValue(c._expression, None, None, ignore_nulls), _emit_ast=False)
89308931
ans._ast = ast
89318932
return ans
89328933

89338934

8935+
@publicapi
8936+
def nth_value(
8937+
e: ColumnOrName, n: int, ignore_nulls: bool = False, _emit_ast: bool = True
8938+
) -> Column:
8939+
"""
8940+
Returns the nth value within an ordered group of values.
8941+
8942+
Example::
8943+
8944+
>>> from snowflake.snowpark.window import Window
8945+
>>> window = Window.partition_by("column1").order_by("column2")
8946+
>>> df = session.create_dataframe([[1, 10], [1, 11], [2, 20], [2, 21]], schema=["column1", "column2"])
8947+
>>> df.select(df["column1"], df["column2"], nth_value(df["column2"], 2).over(window).as_("column2_2nd")).collect()
8948+
[Row(COLUMN1=1, COLUMN2=10, COLUMN2_2ND=11), Row(COLUMN1=1, COLUMN2=11, COLUMN2_2ND=11), Row(COLUMN1=2, COLUMN2=20, COLUMN2_2ND=21), Row(COLUMN1=2, COLUMN2=21, COLUMN2_2ND=21)]
8949+
"""
8950+
# AST.
8951+
ast = None
8952+
if _emit_ast:
8953+
ast = proto.Expr()
8954+
build_builtin_fn_apply(ast, "nth_value", e, n, ignore_nulls)
8955+
8956+
c = _to_col_if_str(e, "nth_value")
8957+
8958+
ans = Column(NthValue(c._expression, n, None, ignore_nulls), _emit_ast=False)
8959+
ans._ast = ast
8960+
return ans
8961+
8962+
89348963
@publicapi
89358964
def ntile(e: Union[int, ColumnOrName], _emit_ast: bool = True) -> Column:
89368965
"""

‎tests/ast/data/functions2.test‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,8 @@ df313 = df.select(log2("A"))
320320

321321
df314 = df.select(instr("A", "test_str"))
322322

323+
df315 = df.select(nth_value("A", 2), nth_value("A", 2, True), nth_value(col("B"), 2, False))
324+
323325
## EXPECTED UNPARSER OUTPUT
324326

325327
df = session.table("table1")
@@ -640,6 +642,8 @@ df313 = df.select(log2("A"))
640642

641643
df314 = df.select(instr("A", "test_str"))
642644

645+
df315 = df.select(nth_value("A", 2, False), nth_value("A", 2, True), nth_value(col("B"), 2, False))
646+
643647
## EXPECTED ENCODED AST
644648

645649
interned_value_table {
@@ -26043,6 +26047,230 @@ body {
2604326047
}
2604426048
}
2604526049
}
26050+
body {
26051+
assign {
26052+
expr {
26053+
sp_dataframe_select__columns {
26054+
cols {
26055+
apply_expr {
26056+
fn {
26057+
builtin_fn {
26058+
name {
26059+
name {
26060+
sp_name_flat {
26061+
name: "nth_value"
26062+
}
26063+
}
26064+
}
26065+
}
26066+
}
26067+
pos_args {
26068+
string_val {
26069+
src {
26070+
end_column: 43
26071+
end_line: 345
26072+
file: 2
26073+
start_column: 26
26074+
start_line: 345
26075+
}
26076+
v: "A"
26077+
}
26078+
}
26079+
pos_args {
26080+
int64_val {
26081+
src {
26082+
end_column: 43
26083+
end_line: 345
26084+
file: 2
26085+
start_column: 26
26086+
start_line: 345
26087+
}
26088+
v: 2
26089+
}
26090+
}
26091+
pos_args {
26092+
bool_val {
26093+
src {
26094+
end_column: 43
26095+
end_line: 345
26096+
file: 2
26097+
start_column: 26
26098+
start_line: 345
26099+
}
26100+
}
26101+
}
26102+
src {
26103+
end_column: 43
26104+
end_line: 345
26105+
file: 2
26106+
start_column: 26
26107+
start_line: 345
26108+
}
26109+
}
26110+
}
26111+
cols {
26112+
apply_expr {
26113+
fn {
26114+
builtin_fn {
26115+
name {
26116+
name {
26117+
sp_name_flat {
26118+
name: "nth_value"
26119+
}
26120+
}
26121+
}
26122+
}
26123+
}
26124+
pos_args {
26125+
string_val {
26126+
src {
26127+
end_column: 68
26128+
end_line: 345
26129+
file: 2
26130+
start_column: 45
26131+
start_line: 345
26132+
}
26133+
v: "A"
26134+
}
26135+
}
26136+
pos_args {
26137+
int64_val {
26138+
src {
26139+
end_column: 68
26140+
end_line: 345
26141+
file: 2
26142+
start_column: 45
26143+
start_line: 345
26144+
}
26145+
v: 2
26146+
}
26147+
}
26148+
pos_args {
26149+
bool_val {
26150+
src {
26151+
end_column: 68
26152+
end_line: 345
26153+
file: 2
26154+
start_column: 45
26155+
start_line: 345
26156+
}
26157+
v: true
26158+
}
26159+
}
26160+
src {
26161+
end_column: 68
26162+
end_line: 345
26163+
file: 2
26164+
start_column: 45
26165+
start_line: 345
26166+
}
26167+
}
26168+
}
26169+
cols {
26170+
apply_expr {
26171+
fn {
26172+
builtin_fn {
26173+
name {
26174+
name {
26175+
sp_name_flat {
26176+
name: "nth_value"
26177+
}
26178+
}
26179+
}
26180+
}
26181+
}
26182+
pos_args {
26183+
apply_expr {
26184+
fn {
26185+
builtin_fn {
26186+
name {
26187+
name {
26188+
sp_name_flat {
26189+
name: "col"
26190+
}
26191+
}
26192+
}
26193+
}
26194+
}
26195+
pos_args {
26196+
string_val {
26197+
src {
26198+
end_column: 88
26199+
end_line: 345
26200+
file: 2
26201+
start_column: 80
26202+
start_line: 345
26203+
}
26204+
v: "B"
26205+
}
26206+
}
26207+
src {
26208+
end_column: 88
26209+
end_line: 345
26210+
file: 2
26211+
start_column: 80
26212+
start_line: 345
26213+
}
26214+
}
26215+
}
26216+
pos_args {
26217+
int64_val {
26218+
src {
26219+
end_column: 99
26220+
end_line: 345
26221+
file: 2
26222+
start_column: 70
26223+
start_line: 345
26224+
}
26225+
v: 2
26226+
}
26227+
}
26228+
pos_args {
26229+
bool_val {
26230+
src {
26231+
end_column: 99
26232+
end_line: 345
26233+
file: 2
26234+
start_column: 70
26235+
start_line: 345
26236+
}
26237+
}
26238+
}
26239+
src {
26240+
end_column: 99
26241+
end_line: 345
26242+
file: 2
26243+
start_column: 70
26244+
start_line: 345
26245+
}
26246+
}
26247+
}
26248+
df {
26249+
sp_dataframe_ref {
26250+
id {
26251+
bitfield1: 1
26252+
}
26253+
}
26254+
}
26255+
src {
26256+
end_column: 100
26257+
end_line: 345
26258+
file: 2
26259+
start_column: 16
26260+
start_line: 345
26261+
}
26262+
variadic: true
26263+
}
26264+
}
26265+
symbol {
26266+
value: "df315"
26267+
}
26268+
uid: 160
26269+
var_id {
26270+
bitfield1: 160
26271+
}
26272+
}
26273+
}
2604626274
client_ast_version: 1
2604726275
client_language {
2604826276
python_language {

0 commit comments

Comments
 (0)