diff --git a/rust/lance-graph/src/datafusion_planner.rs b/rust/lance-graph/src/datafusion_planner.rs index 98ee9130f..9ecdd4616 100644 --- a/rust/lance-graph/src/datafusion_planner.rs +++ b/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) diff --git a/rust/lance-graph/tests/test_datafusion_pipeline.rs b/rust/lance-graph/tests/test_datafusion_pipeline.rs index 9acc82619..3cd5b1e74 100644 --- a/rust/lance-graph/tests/test_datafusion_pipeline.rs +++ b/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::() + .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::() + .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::() + .unwrap(); + + let sum_col = result + .column_by_name("total_age") + .unwrap() + .as_any() + .downcast_ref::() + .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() + ); +}