diff --git a/src/daft-connect/src/translation/expr/unresolved_function.rs b/src/daft-connect/src/translation/expr/unresolved_function.rs index 230924c5de..10e39e7047 100644 --- a/src/daft-connect/src/translation/expr/unresolved_function.rs +++ b/src/daft-connect/src/translation/expr/unresolved_function.rs @@ -24,12 +24,36 @@ pub fn unresolved_to_daft_expr(f: &UnresolvedFunction) -> eyre::Result handle_count(arguments).wrap_err("Failed to handle count function"), + "<" => handle_binary_op(arguments, daft_dsl::Operator::Lt) + .wrap_err("Failed to handle < function"), + ">" => handle_binary_op(arguments, daft_dsl::Operator::Gt) + .wrap_err("Failed to handle > function"), + "<=" => handle_binary_op(arguments, daft_dsl::Operator::LtEq) + .wrap_err("Failed to handle <= function"), + ">=" => handle_binary_op(arguments, daft_dsl::Operator::GtEq) + .wrap_err("Failed to handle >= function"), "isnotnull" => handle_isnotnull(arguments).wrap_err("Failed to handle isnotnull function"), "isnull" => handle_isnull(arguments).wrap_err("Failed to handle isnull function"), n => bail!("Unresolved function {n} not yet supported"), } } +pub fn handle_binary_op( + arguments: Vec, + op: daft_dsl::Operator, +) -> eyre::Result { + let arguments: [daft_dsl::ExprRef; 2] = match arguments.try_into() { + Ok(arguments) => arguments, + Err(arguments) => { + bail!("requires exactly two arguments; got {arguments:?}"); + } + }; + + let [left, right] = arguments; + + Ok(daft_dsl::binary_op(op, left, right)) +} + pub fn handle_count(arguments: Vec) -> eyre::Result { let arguments: [daft_dsl::ExprRef; 1] = match arguments.try_into() { Ok(arguments) => arguments,