From 81d6894eacba5c1d7b9c05ce2b68947db7f2cd5c 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   | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/src/daft-connect/src/translation/expr/unresolved_function.rs b/src/daft-connect/src/translation/expr/unresolved_function.rs
index fc8e7f39b1..7e3c42552b 100644
--- a/src/daft-connect/src/translation/expr/unresolved_function.rs
+++ b/src/daft-connect/src/translation/expr/unresolved_function.rs
@@ -1,5 +1,4 @@
 use daft_core::count_mode::CountMode;
-use daft_schema::dtype::DataType;
 use eyre::{bail, Context};
 use spark_connect::expression::UnresolvedFunction;
 
@@ -25,10 +24,28 @@ 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"),
         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,