Another attempt at an astable flag

Project.toml

@@ -6,14 +6,15 @@ version = "0.9.1"
 Chain = "8be319e6-bccf-4806-a6f7-6fae938471bc"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [compat]
+Chain = "0.4"
 DataFrames = "1"
 MacroTools = "0.5"
 Reexport = "0.2, 1"
 julia = "1"
-Chain = "0.4"
 
 [extras]
 CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"

docs/src/index.md

@@ -22,6 +22,7 @@ In addition, DataFramesMeta provides
   convenient syntax. 
 * `@byrow` for applying functions to each row of a data frame (only supported inside other macros).
 * `@passmissing` for propagating missing values inside row-wise DataFramesMeta.jl transformations.
+* `@astable` to create multiple columns within a single transformation.
 * `@chain`, from [Chain.jl](https://github.com/jkrumbiegel/Chain.jl) for piping the above macros together, similar to [magrittr](https://cran.r-project.org/web/packages/magrittr/vignettes/magrittr.html)'s
   `%>%` in R. 
 
@@ -396,11 +397,37 @@ julia> @rtransform df @passmissing x = parse(Int, :x_str)
    3 │ missing  missing
 ```
 
+## Creating multiple columns at once with `@astable`
+
+Often new variables may depend on the same intermediate calculations. `@astable` makes it easy to create multiple
+new variables in the same operation, yet have them share
+information. 
+
+In a single block, all assignments of the form `:y = f(:x)` 
+or `$y = f(:x)` at the top-level are generate new columns.
+
+```
+julia> df = DataFrame(a = [1, 2, 3], b = [400, 500, 600]);
+
+julia> @transform df @astable begin 
+           ex = extrema(:b)
+           :b_first = :b .- first(ex)
+           :b_last = :b .- last(ex)
+       end
+3×4 DataFrame
+ Row │ a      b      b_first  b_last 
+     │ Int64  Int64  Int64    Int64  
+─────┼───────────────────────────────
+   1 │     1    400        0    -200
+   2 │     2    500      100    -100
+   3 │     3    600      200       0
+```
+
+
 ## [Working with column names programmatically with `$`](@id dollar)
 
 DataFramesMeta provides the special syntax `$` for referring to 
-columns in a data frame via a `Symbol`, string, or column position as either
-a literal or a variable. 
+columns in a data frame via a `Symbol`, string, or column position as either a literal or a variable. 
 
 ```julia
 df = DataFrame(A = 1:3, B = [2, 1, 2])

src/DataFramesMeta.jl

@@ -4,6 +4,8 @@ using Reexport
 
 using MacroTools
 
+using OrderedCollections: OrderedCollections
+
 @reexport using DataFrames
 
 @reexport using Chain
@@ -16,12 +18,13 @@ export @with,
        @transform, @select, @transform!, @select!,
        @rtransform, @rselect, @rtransform!, @rselect!,
        @eachrow, @eachrow!,
-       @byrow, @passmissing,
+       @byrow, @passmissing, @astable,
        @based_on, @where # deprecated
 
 const DOLLAR = raw"$"
 
 include("parsing.jl")
+include("parsing_astable.jl")
 include("macros.jl")
 include("linqmacro.jl")
 include("eachrow.jl")

src/macros.jl

@@ -350,6 +350,120 @@ macro passmissing(args...)
     throw(ArgumentError("@passmissing only works inside DataFramesMeta macros."))
 end
 
+"""
+    astable(args...)
+
+Return a `NamedTuple` from a single transformation inside DataFramesMeta.jl macros.
+
+`@astable` acts on a single block. It works through all top-level expressions
+and collects all such expressions of the form `:y = ...`, i.e. assignments to a
+`Symbol`, which is a syntax error outside of DataFramesMeta.jl macros. At the end of the
+expression, all assignments are collected into a `NamedTuple` to be used
+with the `AsTable` destination in the DataFrames.jl transformation
+mini-language.
+
+Concretely, the expressions
+
+```
+df = DataFrame(a = 1)
+
+@rtransform df @astable begin
+    :x = 1
+    y = 50
+    :z = :x + y + :a
+end
+```
+
+become the pair
+
+```
+function f(a)
+    x_t = 1
+    y = 50
+    z_t = x_t + y + a
+
+    (; x = x_t, z = z_t)
+end
+
+transform(df, [:a] => ByRow(f) => AsTable)
+```
+
+`@astable` has two major advantages at the cost of increasing complexity.
+First, `@astable` makes it easy to create multiple columns from a single
+transformation, which share a scope. For example, `@astable` allows
+for the following
+
+```
+@transform df @astable begin
+    m = mean(:x)
+    :x_demeaned = :x .- m
+    :x2_demeaned = :x2 .- m
+end
+```
+
+The creation of `:x_demeaned` and `:x2_demeaned` both share the variable `m`,
+which does not need to be calculated twice.
+
+Second, `@astable` is useful when performing intermediate calculations
+and storing their results in new columns. For example, the following fails.
+
+```
+@rtransform df begin
+    :new_col_1 = :x + :y
+    :new_col_2 = :new_col_1 + :z
+end
+```
+
+This because DataFrames.jl does not guarantee sequential evaluation of
+transformations. `@astable` solves this problem
+
+@rtransform df @astable begin
+    :new_col_1 = :x + :y
+    :new_col_2 = :new_col_1 + :z
+end
+
+Column assignment in `@astable` follows the same rules as
+column assignment more generally. Construct a new column
+from a string by escaping it with `$DOLLAR`.
+
+### Examples
+
+```
+julia> df = DataFrame(a = [1, 2, 3], b = [4, 5, 6]);
+
+julia> d = @rtransform df @astable begin
+           :x = 1
+           y = 5
+           :z = :x + y
+       end
+3×4 DataFrame
+ Row │ a      b      x      z
+     │ Int64  Int64  Int64  Int64
+─────┼────────────────────────────
+   1 │     1      4      1      6
+   2 │     2      5      1      6
+   3 │     3      6      1      6
+
+julia> df = DataFrame(a = [1, 1, 2, 2], b = [5, 6, 70, 80]);
+
+julia> @by df :a @astable begin
+            ex = extrema(:b)
+            :min_b = first(ex)
+            :max_b = last(ex)
+       end
+2×3 DataFrame
+ Row │ a      min_b  max_b
+     │ Int64  Int64  Int64
+─────┼─────────────────────
+   1 │     1      5      6
+   2 │     2     70     80
+```
+
+"""
+macro astable(args...)
+    throw(ArgumentError("@astable only works inside DataFramesMeta macros."))
+end
+
 ##############################################################################
 ##
 ## @with
@@ -1546,17 +1660,6 @@ function combine_helper(x, args...; deprecation_warning = false)
 
     exprs, outer_flags = create_args_vector(args...)
 
-    fe = first(exprs)
-    if length(exprs) == 1 &&
-        get_column_expr(fe) === nothing &&
-        !(fe.head == :(=) || fe.head == :kw)
-
-        @warn "Returning a Table object from @by and @combine now requires `$(DOLLAR)AsTable` on the LHS."
-
-        lhs = Expr(:$, :AsTable)
-        exprs = ((:($lhs = $fe)),)
-    end
-
     t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
 
     quote
@@ -1666,16 +1769,6 @@ end
 function by_helper(x, what, args...)
     # Only allow one argument when returning a Table object
     exprs, outer_flags = create_args_vector(args...)
-    fe = first(exprs)
-    if length(exprs) == 1 &&
-        get_column_expr(fe) === nothing &&
-        !(fe.head == :(=) || fe.head == :kw)
-
-        @warn "Returning a Table object from @by and @combine now requires `\$AsTable` on the LHS."
-
-        lhs = Expr(:$, :AsTable)
-        exprs = ((:($lhs = $fe)),)
-    end
 
     t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
 

src/parsing.jl

@@ -91,7 +91,8 @@ is_macro_head(ex::Expr, name) = ex.head == :macrocall && ex.args[1] == Symbol(na
 
 const BYROW_SYM = Symbol("@byrow")
 const PASSMISSING_SYM = Symbol("@passmissing")
-const DEFAULT_FLAGS = (;BYROW_SYM => Ref(false), PASSMISSING_SYM => Ref(false))
+const ASTABLE_SYM = Symbol("@astable")
+const DEFAULT_FLAGS = (;BYROW_SYM => Ref(false), PASSMISSING_SYM => Ref(false), ASTABLE_SYM => Ref(false))
 
 extract_macro_flags(ex, exprflags = deepcopy(DEFAULT_FLAGS)) = (ex, exprflags)
 function extract_macro_flags(ex::Expr, exprflags = deepcopy(DEFAULT_FLAGS))
@@ -269,7 +270,13 @@ function fun_to_vec(ex::Expr;
         return ex_col
     end
 
-    if no_dest
+    if final_flags[ASTABLE_SYM][]
+        src, fun = get_source_fun_astable(ex; exprflags = final_flags)
+
+        return :($src => $fun => AsTable)
+    end
+
+    if no_dest # subset and with
         src, fun = get_source_fun(ex, exprflags = final_flags)
         return quote
             $src => $fun
@@ -359,7 +366,7 @@ function create_args_vector(arg; wrap_byrow::Bool=false)
         outer_flags[BYROW_SYM][] = true
     end
 
-    if arg isa Expr && arg.head == :block
+    if arg isa Expr && arg.head == :block && !outer_flags[ASTABLE_SYM][]
         x = MacroTools.rmlines(arg).args
     else
         x = Any[arg]

src/parsing_astable.jl

@@ -0,0 +1,102 @@
+function conditionally_add_symbols!(inputs_to_function::AbstractDict,
+                                    lhs_assignments::OrderedCollections.OrderedDict, col)
+    # if it's already been assigned at top-level,
+    # don't add it to the inputs
+    if haskey(lhs_assignments, col)
+        return lhs_assignments[col]
+    else
+        return addkey!(inputs_to_function, col)
+    end
+end
+
+replace_syms_astable!(inputs_to_function::AbstractDict,
+                      lhs_assignments::OrderedCollections.OrderedDict, x) = x
+replace_syms_astable!(inputs_to_function::AbstractDict,
+                      lhs_assignments::OrderedCollections.OrderedDict, q::QuoteNode) =
+    conditionally_add_symbols!(inputs_to_function, lhs_assignments, q)
+
+function replace_syms_astable!(inputs_to_function::AbstractDict,
+                               lhs_assignments::OrderedCollections.OrderedDict, e::Expr)
+    if onearg(e, :^)
+        return e.args[2]
+    end
+
+    col = get_column_expr(e)
+    if col !== nothing
+        return conditionally_add_symbols!(inputs_to_function, lhs_assignments, col)
+    elseif e.head == :.
+        return replace_dotted_astable!(inputs_to_function, lhs_assignments, e)
+    else
+        return mapexpr(x -> replace_syms_astable!(inputs_to_function, lhs_assignments, x), e)
+    end
+end
+
+protect_replace_syms_astable!(inputs_to_function::AbstractDict,
+                              lhs_assignments::OrderedCollections.OrderedDict, e) = e
+protect_replace_syms_astable!(inputs_to_function::AbstractDict,
+                              lhs_assignments::OrderedCollections.OrderedDict, e::Expr) =
+    replace_syms!(inputs_to_function, lhs_assignments, e)
+
+function replace_dotted_astable!(inputs_to_function::AbstractDict,
+                                 lhs_assignments::OrderedCollections.OrderedDict, e)
+    x_new = replace_syms_astable!(inputs_to_function, lhs_assignments, e.args[1])
+    y_new = protect_replace_syms_astable!(inputs_to_function, lhs_assignments, e.args[2])
+    Expr(:., x_new, y_new)
+end
+
+is_column_assigment(ex) = false
+function is_column_assigment(ex::Expr)
+    ex.head == :(=) && (get_column_expr(ex.args[1]) !== nothing)
+end
+
+# Taken from MacroTools.jl
+# No docstring so assumed unstable
+block(ex) = isexpr(ex, :block) ? ex : :($ex;)
+
+function get_source_fun_astable(ex; exprflags = deepcopy(DEFAULT_FLAGS))
+    inputs_to_function = Dict{Any, Symbol}()
+    lhs_assignments = OrderedCollections.OrderedDict{Any, Symbol}()
+
+    # Make sure all top-level assignments are
+    # in the args vector
+    ex = block(MacroTools.flatten(ex))
+    exprs = map(ex.args) do arg
+        if is_column_assigment(arg)
+            lhs = get_column_expr(arg.args[1])
+            rhs = arg.args[2]
+            new_ex = replace_syms_astable!(inputs_to_function, lhs_assignments, arg.args[2])
+            if haskey(inputs_to_function, lhs)
+                new_lhs = inputs_to_function[lhs]
+                lhs_assignments[lhs] = new_lhs
+            else
+                new_lhs = addkey!(lhs_assignments, lhs)
+            end
+
+            Expr(:(=), new_lhs, new_ex)
+        else
+            replace_syms_astable!(inputs_to_function, lhs_assignments, arg)
+        end
+    end
+    source = :(DataFramesMeta.make_source_concrete($(Expr(:vect, keys(inputs_to_function)...))))
+
+    inputargs = Expr(:tuple, values(inputs_to_function)...)
+    nt_iterator = (:(Symbol($k) => $v) for (k, v) in lhs_assignments)
+    nt_expr = Expr(:tuple, Expr(:parameters, nt_iterator...))
+    body = Expr(:block, Expr(:block, exprs...), nt_expr)
+
+    fun = quote
+        $inputargs -> begin
+            $body
+        end
+    end
+
+    # TODO: Add passmissing support by
+    # checking if any input arguments missing,
+    # and if-so, making a named tuple with
+    # missing values
+    if exprflags[BYROW_SYM][]
+        fun = :(ByRow($fun))
+    end
+
+    return source, fun
+end