Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 11 additions & 0 deletions rust/lance-graph/src/datafusion_planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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)
Expand Down
136 changes: 136 additions & 0 deletions rust/lance-graph/tests/test_datafusion_pipeline.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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()
);
}
Loading