Skip to content

Commit 4d749c7

Browse files
timsaucerclaude
andcommitted
Add missing conditional functions: greatest, least, nvl2, ifnull (#1449)
Expose four conditional functions from upstream DataFusion that were not yet available in the Python bindings. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent be8dd9d commit 4d749c7

2 files changed

Lines changed: 89 additions & 0 deletions

File tree

crates/core/src/functions.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -494,6 +494,8 @@ expr_fn!(length, string);
494494
expr_fn!(char_length, string);
495495
expr_fn!(chr, arg, "Returns the character with the given code.");
496496
expr_fn_vec!(coalesce);
497+
expr_fn_vec!(greatest);
498+
expr_fn_vec!(least);
497499
expr_fn!(cos, num);
498500
expr_fn!(cosh, num);
499501
expr_fn!(cot, num);
@@ -543,6 +545,11 @@ expr_fn!(
543545
x y,
544546
"Returns x if x is not NULL otherwise returns y."
545547
);
548+
expr_fn!(
549+
nvl2,
550+
x y z,
551+
"Returns y if x is not NULL; otherwise returns z."
552+
);
546553
expr_fn!(nullif, arg_1 arg_2);
547554
expr_fn!(
548555
octet_length,
@@ -983,13 +990,15 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
983990
m.add_wrapped(wrap_pyfunction!(floor))?;
984991
m.add_wrapped(wrap_pyfunction!(from_unixtime))?;
985992
m.add_wrapped(wrap_pyfunction!(gcd))?;
993+
m.add_wrapped(wrap_pyfunction!(greatest))?;
986994
// m.add_wrapped(wrap_pyfunction!(grouping))?;
987995
m.add_wrapped(wrap_pyfunction!(in_list))?;
988996
m.add_wrapped(wrap_pyfunction!(initcap))?;
989997
m.add_wrapped(wrap_pyfunction!(isnan))?;
990998
m.add_wrapped(wrap_pyfunction!(iszero))?;
991999
m.add_wrapped(wrap_pyfunction!(levenshtein))?;
9921000
m.add_wrapped(wrap_pyfunction!(lcm))?;
1001+
m.add_wrapped(wrap_pyfunction!(least))?;
9931002
m.add_wrapped(wrap_pyfunction!(left))?;
9941003
m.add_wrapped(wrap_pyfunction!(length))?;
9951004
m.add_wrapped(wrap_pyfunction!(ln))?;
@@ -1007,6 +1016,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
10071016
m.add_wrapped(wrap_pyfunction!(named_struct))?;
10081017
m.add_wrapped(wrap_pyfunction!(nanvl))?;
10091018
m.add_wrapped(wrap_pyfunction!(nvl))?;
1019+
m.add_wrapped(wrap_pyfunction!(nvl2))?;
10101020
m.add_wrapped(wrap_pyfunction!(now))?;
10111021
m.add_wrapped(wrap_pyfunction!(nullif))?;
10121022
m.add_wrapped(wrap_pyfunction!(octet_length))?;

python/datafusion/functions.py

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,8 @@
151151
"floor",
152152
"from_unixtime",
153153
"gcd",
154+
"greatest",
155+
"ifnull",
154156
"in_list",
155157
"initcap",
156158
"isnan",
@@ -159,6 +161,7 @@
159161
"last_value",
160162
"lcm",
161163
"lead",
164+
"least",
162165
"left",
163166
"length",
164167
"levenshtein",
@@ -215,6 +218,7 @@
215218
"ntile",
216219
"nullif",
217220
"nvl",
221+
"nvl2",
218222
"octet_length",
219223
"order_by",
220224
"overlay",
@@ -1030,6 +1034,44 @@ def gcd(x: Expr, y: Expr) -> Expr:
10301034
return Expr(f.gcd(x.expr, y.expr))
10311035

10321036

1037+
def greatest(*args: Expr) -> Expr:
1038+
"""Returns the greatest value from a list of expressions.
1039+
1040+
Returns NULL if all expressions are NULL.
1041+
1042+
Examples:
1043+
>>> ctx = dfn.SessionContext()
1044+
>>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]})
1045+
>>> result = df.select(
1046+
... dfn.functions.greatest(dfn.col("a"), dfn.col("b")).alias("greatest"))
1047+
>>> result.collect_column("greatest")[0].as_py()
1048+
2
1049+
>>> result.collect_column("greatest")[1].as_py()
1050+
3
1051+
"""
1052+
args = [arg.expr for arg in args]
1053+
return Expr(f.greatest(*args))
1054+
1055+
1056+
def ifnull(x: Expr, y: Expr) -> Expr:
1057+
"""Returns ``x`` if ``x`` is not NULL. Otherwise returns ``y``.
1058+
1059+
This is an alias for :py:func:`nvl`.
1060+
1061+
Examples:
1062+
>>> ctx = dfn.SessionContext()
1063+
>>> df = ctx.from_pydict({"a": [None, 1], "b": [0, 0]})
1064+
>>> result = df.select(
1065+
... dfn.functions.ifnull(dfn.col("a"), dfn.col("b")).alias("ifnull")
1066+
... )
1067+
>>> result.collect_column("ifnull")[0].as_py()
1068+
0
1069+
>>> result.collect_column("ifnull")[1].as_py()
1070+
1
1071+
"""
1072+
return nvl(x, y)
1073+
1074+
10331075
def initcap(string: Expr) -> Expr:
10341076
"""Set the initial letter of each word to capital.
10351077
@@ -1083,6 +1125,25 @@ def lcm(x: Expr, y: Expr) -> Expr:
10831125
return Expr(f.lcm(x.expr, y.expr))
10841126

10851127

1128+
def least(*args: Expr) -> Expr:
1129+
"""Returns the least value from a list of expressions.
1130+
1131+
Returns NULL if all expressions are NULL.
1132+
1133+
Examples:
1134+
>>> ctx = dfn.SessionContext()
1135+
>>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]})
1136+
>>> result = df.select(
1137+
... dfn.functions.least(dfn.col("a"), dfn.col("b")).alias("least"))
1138+
>>> result.collect_column("least")[0].as_py()
1139+
1
1140+
>>> result.collect_column("least")[1].as_py()
1141+
1
1142+
"""
1143+
args = [arg.expr for arg in args]
1144+
return Expr(f.least(*args))
1145+
1146+
10861147
def left(string: Expr, n: Expr) -> Expr:
10871148
"""Returns the first ``n`` characters in the ``string``.
10881149
@@ -1267,6 +1328,24 @@ def nvl(x: Expr, y: Expr) -> Expr:
12671328
return Expr(f.nvl(x.expr, y.expr))
12681329

12691330

1331+
def nvl2(x: Expr, y: Expr, z: Expr) -> Expr:
1332+
"""Returns ``y`` if ``x`` is not NULL. Otherwise returns ``z``.
1333+
1334+
Examples:
1335+
>>> ctx = dfn.SessionContext()
1336+
>>> df = ctx.from_pydict({"a": [None, 1], "b": [10, 20], "c": [30, 40]})
1337+
>>> result = df.select(
1338+
... dfn.functions.nvl2(
1339+
... dfn.col("a"), dfn.col("b"), dfn.col("c")).alias("nvl2")
1340+
... )
1341+
>>> result.collect_column("nvl2")[0].as_py()
1342+
30
1343+
>>> result.collect_column("nvl2")[1].as_py()
1344+
20
1345+
"""
1346+
return Expr(f.nvl2(x.expr, y.expr, z.expr))
1347+
1348+
12701349
def octet_length(arg: Expr) -> Expr:
12711350
"""Returns the number of bytes of a string.
12721351

0 commit comments

Comments
 (0)