Skip to content

Commit

Permalink
adds sql functions
Browse files Browse the repository at this point in the history
  • Loading branch information
SeanTroyUWO committed Sep 15, 2023
1 parent cd9969c commit c7c6e4e
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 1 deletion.
45 changes: 44 additions & 1 deletion crates/polars-sql/src/functions.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use polars_core::prelude::{polars_bail, polars_err, PolarsError, PolarsResult};
use polars_lazy::dsl::Expr;
use polars_plan::dsl::count;
use polars_plan::dsl::{coalesce, count, when};
use polars_plan::logical_plan::LiteralValue;
use polars_plan::prelude::lit;
use sqlparser::ast::{
Expand Down Expand Up @@ -232,6 +232,16 @@ pub(crate) enum PolarsSqlFunctions {
/// SELECT UPPER(column_1) from df;
/// ```
Upper,
/// SQL 'nullif' function
/// ```sql
/// SELECT NULLIF(column_1, column_2) from df;
/// ```
NullIf,
/// SQL 'coalesce' function
/// ```sql
/// SELECT COALESCE(column_1, ...) from df;
/// ```
Coalesce,

// ----
// Aggregate functions
Expand Down Expand Up @@ -389,6 +399,7 @@ impl PolarsSqlFunctions {
"cbrt",
"ceil",
"ceiling",
"coalesce",
"cos",
"cosd",
"cot",
Expand All @@ -411,6 +422,7 @@ impl PolarsSqlFunctions {
"ltrim",
"max",
"min",
"nullif",
"octet_length",
"pi",
"pow",
Expand Down Expand Up @@ -476,6 +488,12 @@ impl TryFrom<&'_ SQLFunction> for PolarsSqlFunctions {
"cbrt" => Self::Cbrt,
"round" => Self::Round,

// ----
// Comparison functions
// ----
"nullif" => Self::NullIf,
"coalesce" => Self::Coalesce,

// ----
// String functions
// ----
Expand Down Expand Up @@ -580,6 +598,13 @@ impl SqlFunctionVisitor<'_> {
polars_bail!(InvalidOperation:"Invalid number of arguments for Round: {}",function.args.len());
},
},

// ----
// Comparison functions
// ----
NullIf => self.visit_binary(|l, r: Expr| when(l.clone().eq(r)).then(lit(crate::functions::LiteralValue::Null)).otherwise(l)),
Coalesce => self.visit_variadic(coalesce),

// ----
// String functions
// ----
Expand Down Expand Up @@ -787,6 +812,24 @@ impl SqlFunctionVisitor<'_> {
}
}

fn visit_variadic(&self, f: impl Fn(&[Expr]) -> Expr) -> PolarsResult<Expr> {
self.try_visit_variadic(|e| Ok(f(e)))
}

fn try_visit_variadic(&self, f: impl Fn(&[Expr]) -> PolarsResult<Expr>) -> PolarsResult<Expr> {
let function = self.func;
let args = extract_args(function);
let mut expr_args = vec![];
for arg in args {
if let FunctionArgExpr::Expr(sql_expr) = arg {
expr_args.push(parse_sql_expr(sql_expr, self.ctx)?);
} else {
return self.not_supported_error();
};
}
f(&expr_args)
}

// fn visit_ternary<Arg: FromSqlExpr>(
// &self,
// f: impl Fn(Expr, Arg, Arg) -> Expr,
Expand Down
29 changes: 29 additions & 0 deletions py-polars/tests/unit/test_sql.py
Original file line number Diff line number Diff line change
Expand Up @@ -773,6 +773,35 @@ def test_sql_trim(foods_ipc_path: Path) -> None:
}


def test_sql_nullif_coalesce(foods_ipc_path: Path) -> None:
nums = pl.LazyFrame(
{
"x": [1, None, 2, 3, None, 4],
"y": [5, 4, None, 3, None, 2],
"z": [3, 4, None, 3, None, None],
}
)

res = pl.SQLContext(df=nums).execute(
"""
SELECT
COALESCE(x,y,z) as "coal",
NULLIF(x,y) as "nullif x_y",
NULLIF(y,z) as "nullif y_z",
COALESCE(x, NULLIF(y,z)) as "both"
FROM df
""",
eager=True,
)

assert res.to_dict(False) == {
"coal": [1, 4, 2, 3, None, 4],
"nullif x_y": [1, None, 2, None, None, 4],
"nullif y_z": [5, None, None, None, None, 2],
"both": [1, None, 2, 3, None, 4],
}


def test_sql_order_by(foods_ipc_path: Path) -> None:
foods = pl.scan_ipc(foods_ipc_path)
nums = pl.LazyFrame(
Expand Down

0 comments on commit c7c6e4e

Please sign in to comment.