From 29778d6d34709928ce81f7c0299f3d74a1f176fb Mon Sep 17 00:00:00 2001
From: Grieve <grieve54706@gmail.com>
Date: Fri, 13 Dec 2024 15:09:27 +0800
Subject: [PATCH] feat(core): support querying window function

---
 .../src/logical_plan/analyze/model_anlayze.rs  | 18 ++++++++++++++++++
 wren-core/core/src/mdl/mod.rs                  |  1 +
 wren-core/sqllogictest/test_files/model.slt    |  3 +++
 3 files changed, 22 insertions(+)

diff --git a/wren-core/core/src/logical_plan/analyze/model_anlayze.rs b/wren-core/core/src/logical_plan/analyze/model_anlayze.rs
index 4966b5458..294047331 100644
--- a/wren-core/core/src/logical_plan/analyze/model_anlayze.rs
+++ b/wren-core/core/src/logical_plan/analyze/model_anlayze.rs
@@ -223,6 +223,24 @@ impl ModelAnalyzeRule {
                 scope_mut.add_visited_table(subquery_alias.alias.clone());
                 Ok(Transformed::no(plan))
             }
+            LogicalPlan::Window(window) => {
+                let mut scope_mut = scope.borrow_mut();
+                window
+                    .window_expr
+                    .iter()
+                    .fold(HashSet::new(), |mut set, expr| {
+                        Expr::add_column_refs(expr, &mut set);
+                        set
+                    })
+                    .into_iter()
+                    .try_for_each(|col| {
+                        self.collect_required_column(
+                            Expr::Column(col.to_owned()),
+                            &mut scope_mut,
+                        )
+                    })?;
+                Ok(Transformed::no(plan))
+            }
             _ => Ok(Transformed::no(plan)),
         }
     }
diff --git a/wren-core/core/src/mdl/mod.rs b/wren-core/core/src/mdl/mod.rs
index ed1833b94..fa9b8eca4 100644
--- a/wren-core/core/src/mdl/mod.rs
+++ b/wren-core/core/src/mdl/mod.rs
@@ -478,6 +478,7 @@ mod test {
                 "select o_orderkey, count(*) from test.test.orders where orders.o_totalprice > 10 group by 1",
                 "select totalcost from test.test.profile",
                 "select totalcost from profile",
+                "select sum(c_custkey) over (order by c_name) from test.test.customer limit 1",
         // TODO: support calculated without relationship
         //     "select orderkey_plus_custkey from orders",
         ];
diff --git a/wren-core/sqllogictest/test_files/model.slt b/wren-core/sqllogictest/test_files/model.slt
index 70867ef86..1b334b4a8 100644
--- a/wren-core/sqllogictest/test_files/model.slt
+++ b/wren-core/sqllogictest/test_files/model.slt
@@ -31,6 +31,9 @@ select * from wrenai.public."Customers";
 statement ok
 select count(*) from wrenai.public."Order_items";
 
+statement ok
+select sum("Price") over (order by "Product_id") from wrenai.public."Order_items" limit 1;
+
 # check the count of order_items won't be increased by the relationship calculation
 query B
 select cnt1 = cnt2 from (select count(*) as cnt1 from (select "Customer_state" from wrenai.public."Order_items")), (select count(*) as cnt2 from datafusion.public.order_items) limit 1;