From a2e0bebbabd8bfbdeca4e8b36cb5a7f685b4aa3c Mon Sep 17 00:00:00 2001
From: Andrew Gazelka <andrew.gazelka@gmail.com>
Date: Tue, 19 Nov 2024 23:14:52 -0800
Subject: [PATCH] [FEAT] connect: add binary operators

---
 .../translation/expr/unresolved_function.rs   | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

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<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,