Skip to content

Commit 0b6ea95

Browse files
timsaucerclaude
andauthored
Add missing conditional functions (#1464)
* 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> * Add unit tests for greatest, least, nvl2, and ifnull functions Tests cover multiple data types (integers, strings), null handling (all-null, partial-null), multiple arguments, and ifnull/nvl equivalence. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * Use standard alias docstring pattern for ifnull Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> * remove unused df fixture and fix parameter shadowing * Refactor conditional function tests into parametrized test suite Replace separate test functions for coalesce, greatest, least, nvl, nvl2, ifnull with a single parametrized test using a shared fixture. Adds coverage for nvl, nullif (previously untested), datetime and boolean types, literal fallbacks, and variadic calls. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 645d261 commit 0b6ea95

File tree

3 files changed

+319
-49
lines changed

3 files changed

+319
-49
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!(
498500
contains,
499501
string search_str,
@@ -548,6 +550,11 @@ expr_fn!(
548550
x y,
549551
"Returns x if x is not NULL otherwise returns y."
550552
);
553+
expr_fn!(
554+
nvl2,
555+
x y z,
556+
"Returns y if x is not NULL; otherwise returns z."
557+
);
551558
expr_fn!(nullif, arg_1 arg_2);
552559
expr_fn!(
553560
octet_length,
@@ -989,13 +996,15 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
989996
m.add_wrapped(wrap_pyfunction!(floor))?;
990997
m.add_wrapped(wrap_pyfunction!(from_unixtime))?;
991998
m.add_wrapped(wrap_pyfunction!(gcd))?;
999+
m.add_wrapped(wrap_pyfunction!(greatest))?;
9921000
// m.add_wrapped(wrap_pyfunction!(grouping))?;
9931001
m.add_wrapped(wrap_pyfunction!(in_list))?;
9941002
m.add_wrapped(wrap_pyfunction!(initcap))?;
9951003
m.add_wrapped(wrap_pyfunction!(isnan))?;
9961004
m.add_wrapped(wrap_pyfunction!(iszero))?;
9971005
m.add_wrapped(wrap_pyfunction!(levenshtein))?;
9981006
m.add_wrapped(wrap_pyfunction!(lcm))?;
1007+
m.add_wrapped(wrap_pyfunction!(least))?;
9991008
m.add_wrapped(wrap_pyfunction!(left))?;
10001009
m.add_wrapped(wrap_pyfunction!(length))?;
10011010
m.add_wrapped(wrap_pyfunction!(ln))?;
@@ -1013,6 +1022,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
10131022
m.add_wrapped(wrap_pyfunction!(named_struct))?;
10141023
m.add_wrapped(wrap_pyfunction!(nanvl))?;
10151024
m.add_wrapped(wrap_pyfunction!(nvl))?;
1025+
m.add_wrapped(wrap_pyfunction!(nvl2))?;
10161026
m.add_wrapped(wrap_pyfunction!(now))?;
10171027
m.add_wrapped(wrap_pyfunction!(nullif))?;
10181028
m.add_wrapped(wrap_pyfunction!(octet_length))?;

python/datafusion/functions.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,8 @@
152152
"floor",
153153
"from_unixtime",
154154
"gcd",
155+
"greatest",
156+
"ifnull",
155157
"in_list",
156158
"initcap",
157159
"isnan",
@@ -160,6 +162,7 @@
160162
"last_value",
161163
"lcm",
162164
"lead",
165+
"least",
163166
"left",
164167
"length",
165168
"levenshtein",
@@ -216,6 +219,7 @@
216219
"ntile",
217220
"nullif",
218221
"nvl",
222+
"nvl2",
219223
"octet_length",
220224
"order_by",
221225
"overlay",
@@ -1045,6 +1049,34 @@ def gcd(x: Expr, y: Expr) -> Expr:
10451049
return Expr(f.gcd(x.expr, y.expr))
10461050

10471051

1052+
def greatest(*args: Expr) -> Expr:
1053+
"""Returns the greatest value from a list of expressions.
1054+
1055+
Returns NULL if all expressions are NULL.
1056+
1057+
Examples:
1058+
>>> ctx = dfn.SessionContext()
1059+
>>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]})
1060+
>>> result = df.select(
1061+
... dfn.functions.greatest(dfn.col("a"), dfn.col("b")).alias("greatest"))
1062+
>>> result.collect_column("greatest")[0].as_py()
1063+
2
1064+
>>> result.collect_column("greatest")[1].as_py()
1065+
3
1066+
"""
1067+
exprs = [arg.expr for arg in args]
1068+
return Expr(f.greatest(*exprs))
1069+
1070+
1071+
def ifnull(x: Expr, y: Expr) -> Expr:
1072+
"""Returns ``x`` if ``x`` is not NULL. Otherwise returns ``y``.
1073+
1074+
See Also:
1075+
This is an alias for :py:func:`nvl`.
1076+
"""
1077+
return nvl(x, y)
1078+
1079+
10481080
def initcap(string: Expr) -> Expr:
10491081
"""Set the initial letter of each word to capital.
10501082
@@ -1098,6 +1130,25 @@ def lcm(x: Expr, y: Expr) -> Expr:
10981130
return Expr(f.lcm(x.expr, y.expr))
10991131

11001132

1133+
def least(*args: Expr) -> Expr:
1134+
"""Returns the least value from a list of expressions.
1135+
1136+
Returns NULL if all expressions are NULL.
1137+
1138+
Examples:
1139+
>>> ctx = dfn.SessionContext()
1140+
>>> df = ctx.from_pydict({"a": [1, 3], "b": [2, 1]})
1141+
>>> result = df.select(
1142+
... dfn.functions.least(dfn.col("a"), dfn.col("b")).alias("least"))
1143+
>>> result.collect_column("least")[0].as_py()
1144+
1
1145+
>>> result.collect_column("least")[1].as_py()
1146+
1
1147+
"""
1148+
exprs = [arg.expr for arg in args]
1149+
return Expr(f.least(*exprs))
1150+
1151+
11011152
def left(string: Expr, n: Expr) -> Expr:
11021153
"""Returns the first ``n`` characters in the ``string``.
11031154
@@ -1282,6 +1333,24 @@ def nvl(x: Expr, y: Expr) -> Expr:
12821333
return Expr(f.nvl(x.expr, y.expr))
12831334

12841335

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

0 commit comments

Comments
 (0)