Skip to content

Commit 41ff010

Browse files
jja725jianjian.xie
andauthored
feat: implement arithmetic expression computation in DataFusion planner (#68)
* feat: implement arithmetic expression computation in DataFusion planner Implement proper translation of Cypher arithmetic expressions (+, -, *, /, %) to DataFusion BinaryExpr instead of returning a placeholder lit(0). Changes: - Add full arithmetic operator support in to_df_value_expr() - Map ArithmeticOperator variants to DataFusion Operator equivalents - Update tests to verify proper BinaryExpr generation with correct operators * fix format --------- Co-authored-by: jianjian.xie <jianjian.xie@uber.com>
1 parent 0e1270d commit 41ff010

1 file changed

Lines changed: 44 additions & 14 deletions

File tree

rust/lance-graph/src/datafusion_planner/expression.rs

Lines changed: 44 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,27 @@ pub(crate) fn to_df_value_expr(expr: &ValueExpression) -> Expr {
155155
}
156156
}
157157
}
158-
VE::Arithmetic { .. } => lit(0),
158+
VE::Arithmetic {
159+
left,
160+
operator,
161+
right,
162+
} => {
163+
use crate::ast::ArithmeticOperator as AO;
164+
let l = to_df_value_expr(left);
165+
let r = to_df_value_expr(right);
166+
let op = match operator {
167+
AO::Add => Operator::Plus,
168+
AO::Subtract => Operator::Minus,
169+
AO::Multiply => Operator::Multiply,
170+
AO::Divide => Operator::Divide,
171+
AO::Modulo => Operator::Modulo,
172+
};
173+
Expr::BinaryExpr(BinaryExpr {
174+
left: Box::new(l),
175+
op,
176+
right: Box::new(r),
177+
})
178+
}
159179
}
160180
}
161181

@@ -468,12 +488,12 @@ mod tests {
468488

469489
let df_expr = to_df_value_expr(&expr);
470490
let s = format!("{:?}", df_expr);
471-
// TODO: Arithmetic expressions are not yet implemented (line 124 returns lit(0) placeholder)
472-
// Once implemented, this should return a BinaryExpr with Plus operator
473-
// For now, verify it returns the placeholder without panicking
474-
assert_eq!(
475-
s, "Literal(Int32(0), None)",
476-
"Arithmetic currently returns placeholder lit(0)"
491+
// Arithmetic expressions should now return a BinaryExpr with Plus operator
492+
assert!(s.contains("BinaryExpr"), "Should be a BinaryExpr");
493+
assert!(s.contains("Plus"), "Should contain Plus operator");
494+
assert!(
495+
s.contains("p__age"),
496+
"Should contain the left operand column"
477497
);
478498
}
479499

@@ -482,22 +502,32 @@ mod tests {
482502
use crate::ast::ArithmeticOperator;
483503

484504
let operators = vec![
485-
ArithmeticOperator::Add,
486-
ArithmeticOperator::Subtract,
487-
ArithmeticOperator::Multiply,
488-
ArithmeticOperator::Divide,
505+
(ArithmeticOperator::Add, "Plus"),
506+
(ArithmeticOperator::Subtract, "Minus"),
507+
(ArithmeticOperator::Multiply, "Multiply"),
508+
(ArithmeticOperator::Divide, "Divide"),
489509
];
490510

491-
for op in operators {
511+
for (op, expected_op_str) in operators {
492512
let expr = ValueExpression::Arithmetic {
493513
left: Box::new(ValueExpression::Literal(PropertyValue::Integer(10))),
494514
operator: op,
495515
right: Box::new(ValueExpression::Literal(PropertyValue::Integer(2))),
496516
};
497517

498518
let df_expr = to_df_value_expr(&expr);
499-
// Should successfully translate without panicking
500-
let _ = format!("{:?}", df_expr);
519+
let s = format!("{:?}", df_expr);
520+
// Should translate to BinaryExpr with the correct operator
521+
assert!(
522+
s.contains("BinaryExpr"),
523+
"Arithmetic should produce BinaryExpr"
524+
);
525+
assert!(
526+
s.contains(expected_op_str),
527+
"Expected operator {} in expression: {}",
528+
expected_op_str,
529+
s
530+
);
501531
}
502532
}
503533

0 commit comments

Comments
 (0)