Fix multi-column select! #404 (#405)

* add allow multicol * add tests * add fix * Update src/macros.jl Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl> --------- Co-authored-by: Bogumił Kamiński <bkamins@sgh.waw.pl>
JuliaData · Nov 14, 2024 · d93dcbd · d93dcbd · pdeffebach · Nov 15, 2024
1 parent bffeabb
commit d93dcbd
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 3 deletions.
diff --git a/src/macros.jl b/src/macros.jl
@@ -1930,7 +1930,7 @@ end
 function rselect_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = true)
 
-    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select($x, $(t...); $(kw...))
     end
@@ -1983,7 +1983,7 @@ end
 function select!_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = false)
 
-    t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names = false, outer_flags = outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select!($x, $(t...); $(kw...))
     end
@@ -2119,7 +2119,7 @@ end
 function rselect!_helper(x, args...)
     x, exprs, outer_flags, kw = get_df_args_kwargs(x, args...; wrap_byrow = true)
 
-    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags) for ex in exprs)
+    t = (fun_to_vec(ex; gensym_names=false, outer_flags=outer_flags, allow_multicol=true) for ex in exprs)
     quote
         $select!($x, $(t...); $(kw...))
     end

diff --git a/test/multicol.jl b/test/multicol.jl
@@ -22,6 +22,67 @@ df = DataFrame(A = 1, AA = 2, B = 3)
     @test t == DataFrame(AA = 2, B = 3)
 end
 
+@testset "rselect_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    t = @rselect df Not(:A)
+    @test t == DataFrame(AA = 2, B = 3)
+
+    t = @rselect df All()
+    @test t == DataFrame(A = 1, AA = 2, B = 3)
+
+    t = @rselect df Cols(r"A")
+    @test t == DataFrame(A = 1, AA = 2)
+
+    t = @rselect df Between(:AA, :B)
+    @test t == DataFrame(AA = 2, B = 3)
+end
+
+@testset "select!_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Not(:A)
+    @test df == DataFrame(AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df All()
+    @test df == DataFrame(A = 1, AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Cols(r"A")
+    @test df == DataFrame(A = 1, AA = 2)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @select! df Between(:AA, :B)
+    @test df == DataFrame(AA = 2, B = 3)
+end
+
+@testset "rselect!_multi" begin
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Not(:A)
+    @test df == DataFrame(AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df All()
+    @test df == DataFrame(A = 1, AA = 2, B = 3)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Cols(r"A")
+    @test df == DataFrame(A = 1, AA = 2)
+
+    df = DataFrame(A = 1, AA = 2, B = 3)
+
+    @rselect! df Between(:AA, :B)
+    @test df == DataFrame(AA = 2, B = 3)
+end
+
+
 @testset "othermacros_multi" begin
     df = DataFrame(A = 1, AA = 2, B = 3)