feat: support the SUM function in the datafusion planner (#36)

* feat: support the SUM function in the datafusion planner

* doc: update test comments

Author: ChunxuTang

rust/lance-graph/src/datafusion_planner.rs

@@ -22,6 +22,7 @@ use datafusion::logical_expr::{
     col, lit, BinaryExpr, Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Operator,
 };
 use datafusion_functions_aggregate::count::count;
+use datafusion_functions_aggregate::sum::sum;
 use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
@@ -1541,6 +1542,16 @@ impl DataFusionPlanner {
                             lit(0)
                         }
                     }
+                    "sum" => {
+                        if args.len() == 1 {
+                            let arg_expr = Self::to_df_value_expr(&args[0]);
+                            // Use DataFusion's sum helper function
+                            sum(arg_expr)
+                        } else {
+                            // Invalid argument count - return placeholder
+                            lit(0)
+                        }
+                    }
                     _ => {
                         // Unsupported function - return placeholder for now
                         lit(0)

rust/lance-graph/tests/test_datafusion_pipeline.rs

@@ -2610,3 +2610,139 @@ async fn test_count_property_without_alias_has_descriptive_name() {
         result.schema()
     );
 }
+
+#[tokio::test]
+async fn test_sum_property() {
+    let person_batch = create_person_dataset();
+    let config = GraphConfig::builder()
+        .with_node_label("Person", "id")
+        .build()
+        .unwrap();
+
+    let query = CypherQuery::new("MATCH (p:Person) RETURN sum(p.age) AS total_age")
+        .unwrap()
+        .with_config(config);
+
+    let mut datasets = HashMap::new();
+    datasets.insert("Person".to_string(), person_batch);
+
+    let result = query.execute_datafusion(datasets).await.unwrap();
+
+    assert_eq!(result.num_rows(), 1);
+    let sum_col = result
+        .column_by_name("total_age")
+        .unwrap()
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .unwrap();
+    // Sum of ages: 25 + 35 + 30 + 40 + 28 = 158
+    assert_eq!(sum_col.value(0), 158);
+}
+
+#[tokio::test]
+async fn test_sum_with_filter() {
+    let person_batch = create_person_dataset();
+    let config = GraphConfig::builder()
+        .with_node_label("Person", "id")
+        .build()
+        .unwrap();
+
+    let query =
+        CypherQuery::new("MATCH (p:Person) WHERE p.age >= 30 RETURN sum(p.age) AS total_age")
+            .unwrap()
+            .with_config(config);
+
+    let mut datasets = HashMap::new();
+    datasets.insert("Person".to_string(), person_batch);
+
+    let result = query.execute_datafusion(datasets).await.unwrap();
+
+    assert_eq!(result.num_rows(), 1);
+    let sum_col = result
+        .column_by_name("total_age")
+        .unwrap()
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .unwrap();
+    // Sum of ages >= 30: 35 + 30 + 40 = 105
+    assert_eq!(sum_col.value(0), 105);
+}
+
+#[tokio::test]
+async fn test_sum_with_grouping() {
+    let person_batch = create_person_dataset();
+    let config = GraphConfig::builder()
+        .with_node_label("Person", "id")
+        .build()
+        .unwrap();
+
+    let query =
+        CypherQuery::new("MATCH (p:Person) RETURN p.city, sum(p.age) AS total_age ORDER BY p.city")
+            .unwrap()
+            .with_config(config);
+
+    let mut datasets = HashMap::new();
+    datasets.insert("Person".to_string(), person_batch);
+
+    let result = query.execute_datafusion(datasets).await.unwrap();
+
+    // Should have 5 groups: NULL, Chicago, New York, San Francisco, Seattle
+    assert_eq!(result.num_rows(), 5);
+
+    let city_col = result
+        .column_by_name("p.city")
+        .unwrap()
+        .as_any()
+        .downcast_ref::<StringArray>()
+        .unwrap();
+
+    let sum_col = result
+        .column_by_name("total_age")
+        .unwrap()
+        .as_any()
+        .downcast_ref::<Int64Array>()
+        .unwrap();
+
+    // Verify grouping results (ordered by city, NULL comes first)
+    assert!(city_col.is_null(0)); // David: 40 (NULL city)
+    assert_eq!(sum_col.value(0), 40);
+
+    assert_eq!(city_col.value(1), "Chicago"); // Charlie: 30
+    assert_eq!(sum_col.value(1), 30);
+
+    assert_eq!(city_col.value(2), "New York"); // Alice: 25
+    assert_eq!(sum_col.value(2), 25);
+
+    assert_eq!(city_col.value(3), "San Francisco"); // Bob: 35
+    assert_eq!(sum_col.value(3), 35);
+
+    assert_eq!(city_col.value(4), "Seattle"); // Eve: 28
+    assert_eq!(sum_col.value(4), 28);
+}
+
+#[tokio::test]
+async fn test_sum_without_alias_has_descriptive_name() {
+    let person_batch = create_person_dataset();
+    let config = GraphConfig::builder()
+        .with_node_label("Person", "id")
+        .build()
+        .unwrap();
+
+    let query = CypherQuery::new("MATCH (p:Person) RETURN sum(p.age)")
+        .unwrap()
+        .with_config(config);
+
+    let mut datasets = HashMap::new();
+    datasets.insert("Person".to_string(), person_batch);
+
+    let result = query.execute_datafusion(datasets).await.unwrap();
+
+    assert_eq!(result.num_rows(), 1);
+    // Should have column named "sum(p.age)" not "expr"
+    let sum_col = result.column_by_name("sum(p.age)");
+    assert!(
+        sum_col.is_some(),
+        "Expected column named 'sum(p.age)' but schema is: {:?}",
+        result.schema()
+    );
+}