Skip to content

Commit

Permalink
[FEAT] connect: add binary operators
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgazelka committed Nov 27, 2024
1 parent 2c7437d commit fbc5c6e
Showing 1 changed file with 24 additions and 0 deletions.
24 changes: 24 additions & 0 deletions src/daft-connect/src/translation/expr/unresolved_function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,12 +24,36 @@ pub fn unresolved_to_daft_expr(f: &UnresolvedFunction) -> eyre::Result<daft_dsl:

match function_name.as_str() {
"count" => 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<daft_dsl::ExprRef>,
op: daft_dsl::Operator,
) -> eyre::Result<daft_dsl::ExprRef> {
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<daft_dsl::ExprRef>) -> eyre::Result<daft_dsl::ExprRef> {
let arguments: [daft_dsl::ExprRef; 1] = match arguments.try_into() {
Ok(arguments) => arguments,
Expand Down

0 comments on commit fbc5c6e

Please sign in to comment.