fixes window_frame, window_order, cumsum

TidierOrg · Mar 27, 2024 · 6a9215e · 6a9215e
1 parent b9506e2
commit 6a9215e
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -13,7 +13,7 @@ Currently supported macros:
 - `@slice_min`, `@slice_max`, `@slice_sample`
 - `@show_query`
 - `@collect`
-- `@window_order` and `window_frame` are coming soon
+- `@window_order` and `window_frame`
 
 
 Supported helper functions include

diff --git a/src/TBD_macros.jl b/src/TBD_macros.jl
@@ -189,19 +189,13 @@ macro mutate(sqlquery, mutations...)
                 if isa(expr, Expr) && expr.head == :(=) && isa(expr.args[1], Symbol)
                     col_name = string(expr.args[1])
                     col_expr = expr_to_sql(expr.args[2], sq)  # Ensure you have a function that can handle this conversion
-
-                    # Construct window_clause based on grouping and ordering
-                    partition_clause = !isempty(sq.groupBy) && !isempty(sq.windowFrame) ? "PARTITION BY $(sq.groupBy)" : ""
-                    order_clause = !isempty(sq.window_order) ? "ORDER BY $(sq.window_order)" : ""
-                    frame_clause = !isempty(sq.windowFrame) ? "$(sq.windowFrame)" : "ROWS UNBOUNDED PRECEDING"
-                    window_clause = !isempty(order_clause) || !isempty(partition_clause) ? " OVER ($partition_clause $order_clause $frame_clause)" : ""
 
                     if col_name in all_columns
                         # Replace the existing column expression with the mutation
-                        select_expressions[findfirst(==(col_name), select_expressions)] = string(col_expr, window_clause, " AS ", col_name)
+                        select_expressions[findfirst(==(col_name), select_expressions)] = string(col_expr, " AS ", col_name)
                     else
                         # Append the mutation as a new column expression
-                        push!(select_expressions, string(col_expr, window_clause, " AS ", col_name))
+                        push!(select_expressions, string(col_expr, " AS ", col_name))
                         # Update metadata to include this new column
                         push!(sq.metadata, Dict("name" => col_name, "type" => "UNKNOWN", "current_selxn" => 1))
                     end

diff --git a/src/db_parsing.jl b/src/db_parsing.jl
@@ -144,8 +144,9 @@ function expr_to_sql_lite(expr, sq; from_summarize::Bool)
             end
         elseif @capture(x, cumsum(a_))
             if from_summarize
-                return :(SUM($a))
+                error("cumsum is only available through a windowed @mutate")
             else
+               # sq.windowFrame = "ROWS UNBOUNDED PRECEDING "
                 window_clause = construct_window_clause(sq)
                 return  "SUM($(string(a))) $(window_clause)"
             end
@@ -364,4 +365,26 @@ function parse_blocks(exprs...)
       return (MacroTools.rmlines(exprs[1]).args...,)
     end
     return exprs
-  end
+  end
+
+  function construct_window_clause(sq::SQLQuery ; from_cumsum::Bool = false)
+    # Construct the partition clause, considering both groupBy and window_order
+    partition_clause = !isempty(sq.groupBy) ? "PARTITION BY $(sq.groupBy)" : ""
+    if !isempty(sq.window_order)
+        # If there's already a partition clause, append the order clause; otherwise, start with ORDER BY
+        order_clause = !isempty(partition_clause) ? " ORDER BY $(sq.window_order)" : "ORDER BY $(sq.window_order)"
+    else
+        order_clause = ""
+    end
+    if from_cumsum == true
+        frame_clause = "ROWS UNBOUNDED PRECEDING "
+    else 
+        frame_clause = !isempty(sq.windowFrame) ? sq.windowFrame : ""
+    end
+    # Combine partition, order, and frame clauses for the complete window function clause
+    # Ensure to include space only when needed to avoid syntax issues
+    partition_and_order_clause = partition_clause * (!isempty(order_clause) ? " " * order_clause : "")
+    window_clause = (!isempty(partition_clause) || !isempty(order_clause) || !isempty(frame_clause))  ? "OVER ($partition_and_order_clause $frame_clause)" : "OVER ()"
+
+    return window_clause
+end
diff --git a/src/postgresparsing.jl b/src/postgresparsing.jl
@@ -45,17 +45,18 @@ function expr_to_sql_postgres(expr, sq; from_summarize::Bool)
             end
         elseif @capture(x, cumsum(a_))
             if from_summarize
-                return :(SUM($a))
+                error("cumsum is only available through a windowed @mutate")
             else
-                window_clause = construct_window_clause(sq)
-                return  "SUM($(string(a))) $(window_clause)"
+               # sq.windowFrame = "ROWS UNBOUNDED PRECEDING "
+               window_clause = construct_window_clause(sq, from_cumsum = true)
+               return  "SUM($(string(a))) $(window_clause)"
             end
         #stats agg
         elseif @capture(x, std(a_))
             if from_summarize
                 return :(STDDEV_SAMP($a))
             else
-                window_clause = construct_window_clause(sq)
+                window_clause = construct_window_clause(sq, )
                 return  "STDDEV_SAMP($(string(a))) $(window_clause)"
             end
         elseif @capture(x, cor(a_, b_))