Skip to content

Commit c8dc245

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 5be412b commit c8dc245

File tree

2 files changed

+89
-0
lines changed

2 files changed

+89
-0
lines changed

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,
@@ -981,13 +988,15 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
981988
m.add_wrapped(wrap_pyfunction!(floor))?;
982989
m.add_wrapped(wrap_pyfunction!(from_unixtime))?;
983990
m.add_wrapped(wrap_pyfunction!(gcd))?;
991+
m.add_wrapped(wrap_pyfunction!(greatest))?;
984992
// m.add_wrapped(wrap_pyfunction!(grouping))?;
985993
m.add_wrapped(wrap_pyfunction!(in_list))?;
986994
m.add_wrapped(wrap_pyfunction!(initcap))?;
987995
m.add_wrapped(wrap_pyfunction!(isnan))?;
988996
m.add_wrapped(wrap_pyfunction!(iszero))?;
989997
m.add_wrapped(wrap_pyfunction!(levenshtein))?;
990998
m.add_wrapped(wrap_pyfunction!(lcm))?;
999+
m.add_wrapped(wrap_pyfunction!(least))?;
9911000
m.add_wrapped(wrap_pyfunction!(left))?;
9921001
m.add_wrapped(wrap_pyfunction!(length))?;
9931002
m.add_wrapped(wrap_pyfunction!(ln))?;
@@ -1005,6 +1014,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
10051014
m.add_wrapped(wrap_pyfunction!(named_struct))?;
10061015
m.add_wrapped(wrap_pyfunction!(nanvl))?;
10071016
m.add_wrapped(wrap_pyfunction!(nvl))?;
1017+
m.add_wrapped(wrap_pyfunction!(nvl2))?;
10081018
m.add_wrapped(wrap_pyfunction!(now))?;
10091019
m.add_wrapped(wrap_pyfunction!(nullif))?;
10101020
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
@@ -149,6 +149,8 @@
149149
"floor",
150150
"from_unixtime",
151151
"gcd",
152+
"greatest",
153+
"ifnull",
152154
"in_list",
153155
"initcap",
154156
"isnan",
@@ -157,6 +159,7 @@
157159
"last_value",
158160
"lcm",
159161
"lead",
162+
"least",
160163
"left",
161164
"length",
162165
"levenshtein",
@@ -212,6 +215,7 @@
212215
"ntile",
213216
"nullif",
214217
"nvl",
218+
"nvl2",
215219
"octet_length",
216220
"order_by",
217221
"overlay",
@@ -1027,6 +1031,44 @@ def gcd(x: Expr, y: Expr) -> Expr:
10271031
return Expr(f.gcd(x.expr, y.expr))
10281032

10291033

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

10821124

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

12661327

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

0 commit comments

Comments
 (0)