0.4 bugfixes

docs/src/index.md

@@ -86,20 +86,17 @@ See the benchmarking code for [Stata](https://github.com/codedthinking/Kezdi.jl/
 The function can operate on individual elements,
 ```julia
 get_make(text) = split(text, " ")[1]
-@generate Make = Main.get_make(Model)
+@generate Make = get_make(Model)
 ```
 or on the entire column:
 ```julia
-function geometric_mean(x::AbstractVector)
+function geometric_mean(x::Vector)
     n = length(x)
     return exp(sum(log.(x)) / n)
 end
-@collapse geom_NPG = Main.geometric_mean(MPG), by(Cylinders)
+@collapse geom_NPG = geometric_mean(MPG), by(Cylinders)
 ```
 
-!!! tip "Note: `Main.` prefix"
-    If you define a function in your own code, you need to prefix the function name with `Main.` to use it in other commands. To make use of [Automatic vectorization](@ref), make sure to give the function a vector argument type.
-
 ## Commands
 
 ### Setting and inspecting the global DataFrame

goals.md

@@ -266,3 +266,67 @@ An in-place version of `@with!` should do everything in place. This can mean all
    - non-standard evaluation makes it hard to wrap Kezdi.jl code in functions
 6. For loops
    - implement `scalars()` and automatic expansion of locals in context
+
+# 2024-07-12 In-flight debugging session
+```julia
+julia> using Kezdi
+
+julia> module MyModule
+       myfunc(x) = 2x
+       end
+Main.MyModule
+
+julia> df = DataFrame(x = 1:10)
+10×1 DataFrame
+ Row │ x
+     │ Int64
+─────┼───────
+   1 │     1
+   2 │     2
+   3 │     3
+   4 │     4
+   5 │     5
+   6 │     6
+   7 │     7
+   8 │     8
+   9 │     9
+  10 │    10
+
+julia> @with df @generate y = MyModule.myfunc(x)
+10×2 DataFrame
+ Row │ x      y
+     │ Int64  Int64
+─────┼──────────────
+   1 │     1      2
+   2 │     2      4
+   3 │     3      6
+   4 │     4      8
+   5 │     5     10
+   6 │     6     12
+   7 │     7     14
+   8 │     8     16
+   9 │     9     18
+  10 │    10     20
+```
+
+How about aggreator function?
+
+```julia
+julia> module MyModule
+       myfunc(x) = 2x
+       myaggreg(v::Vector) = sum(x.^2)
+       end
+WARNING: replacing module MyModule.
+Main.MyModule
+
+julia> @with df @egen y = MyModule.myaggreg(x)
+┌ Warning: transform!(var"##237", [:x] => (((x,)->(passmissing(MyModule.myaggreg)).(x)) => $(QuoteNode("y"))))
+└ @ Kezdi ~/Tresorit/Mac/code/julia/Kezdi.jl/src/commands.jl:100
+ERROR: MethodError: no method matching myaggreg(::Int64)
+
+Closest candidates are:
+  myaggreg(::Vector)
+   @ Main.MyModule REPL[8]:3
+```
+
+This means it was vectorized at compile time, but it is found at runtime.

src/codegen.jl

@@ -244,9 +244,10 @@
     end
 end
 
+get_dot_parts(ex::Symbol) = [ex]
 function get_dot_parts(ex::Expr)
     is_dot_reference(ex) || error("Expected a dot reference, got $ex")
-    parts = []
+    parts = Symbol[]
     while is_dot_reference(ex)
         push!(parts, ex.args[2].value)
         ex = ex.args[1]
@@ -274,28 +275,15 @@
 isalphanumeric(str::AbstractString) = all(isalphanumeric, str)
 
 isassignment(expr::Any) = expr isa Expr && expr.head == :(=) && length(expr.args) == 2
-function operates_on_vector(expr::Any) 
-    try
-        length(methodswith(Vector, eval(expr); supertypes=true)) > 0
-    catch e
-        if isa(e, UndefVarError)
-            return false
-        else
-            rethrow(e)
-        end
-    end
-end
+operates_on_missing(expr::Any) = (expr isa Symbol && expr == :ismissing) || operates_on_type(expr, Missing)
+operates_on_vector(expr::Any) = operates_on_type(expr, Vector)
 
-function operates_on_missing(expr::Any)
-    expr isa Symbol && expr == :ismissing && return true 
+function operates_on_type(expr::Any, T::Type)
     try
-        length(methodswith(Missing, eval(expr); supertypes=true)) > 0
-    catch e
-        if isa(e, UndefVarError)
-            return false
-        else
-            rethrow(e)
-        end
+        return length(methodswith(T, Main.eval(expr); supertypes=true)) > 0
+    catch ee
+        !isa(ee, UndefVarError) && rethrow(ee)
+        return false
     end
 end
 

src/commands.jl

@@ -44,10 +44,11 @@ function rewrite(::Val{:replace}, command::Command)
         else
             $setup
             eltype_RHS = $RHS isa AbstractVector ? eltype($RHS) : typeof($RHS)
-            if eltype_RHS != eltype($target_df[!, $target_column])
-                local $third_vector = Vector{eltype_RHS}(undef, nrow($local_copy))
+            eltype_LHS = eltype($local_copy[.!$bitmask, $target_column])
+            if eltype_RHS != eltype_LHS
+                local $third_vector = Vector{promote_type(eltype_LHS, eltype_RHS)}(undef, nrow($local_copy))
                 $third_vector[$bitmask] .= $RHS
-                $third_vector[.!$bitmask] .= $local_copy[!, $target_column][.!$bitmask]
+                $third_vector[.!$bitmask] .= $local_copy[.!$bitmask, $target_column]
                 $local_copy[!, $target_column] = $third_vector
             else
                 $target_df[!, $target_column] .= $RHS

test/codegen.jl

@@ -1,3 +1,10 @@
+module MyModule
+myfunc(x) = 2x
+myaggreg(v::Vector) = sum(x.^2)
+mymiss(::Missing) = missing
+mymiss(x) = 3x
+end
+
 @testset "Replace variable references" begin
     @test_expr replace_variable_references(:(x + y + f(z) - g.(x))) == :(:x + :y + f(:z) - g.(:x))
     @test_expr replace_variable_references(:(f(x, <=))) == :(f(:x, <=))
@@ -39,4 +46,10 @@ end
     @testset "Unknown functions are passed through `passmissing`" begin
         @test_expr vectorize_function_calls(:(y = Dates.year(x))) == :(y = (passmissing(Dates.year)).(x))
     end
+    @testset "Functions in other modules" begin
+        using .MyModule
+        @test vectorize_function_calls(:(MyModule.myfunc(x))) == :((passmissing(MyModule.myfunc)).(x))       
+        @test vectorize_function_calls(:(MyModule.myaggreg(x))) == :(MyModule.myaggreg(keep_only_values(x)))  
+        @test vectorize_function_calls(:(MyModule.mymiss(x))) == :(MyModule.mymiss.(x))     
+    end
 end

test/commands.jl

@@ -50,10 +50,9 @@
         df = DataFrame(x=[[1, 2], [3, 4], [5, 6], [7, 8]])
         @test (@with df @generate x1 = getindex(x, 1)).x1 == [1, 3, 5, 7]
         @test (@with df @generate x2 = getindex(x, 2)).x2 == [2, 4, 6, 8]
-    end
-
-    @testset "Error handling" begin
-        @test_throws Exception @with df @generate x = 1
+        df = DataFrame(text = ["a,b", "c,d,e", "f"])
+        df2 = @with df @generate n_terms = length.(split.(text, ","))
+        @test df2.n_terms == [2, 3, 1]
     end
 end
 
@@ -87,6 +86,19 @@ end
         @test eltype(df.x) == eltype(df3.x)
     end
 
+    @testset "Mixed types" begin
+        df = DataFrame(x=[1, 2, 3])
+        @test eltype((@with df @replace x = 1.1 @if _n == 1).x) <: AbstractFloat
+        @test eltype((@with df @replace x = missing @if _n == 1).x) == Union{Missing, Int}
+        @test eltype((@with df @replace x = "a" @if _n == 1).x) == Any
+        df = DataFrame(x=[missing, 2, 3])
+        @test eltype((@with df @replace x = 1 @if _n == 1).x) == Union{Int, Missing}
+        df = DataFrame(x=[1.1, 2, 3])
+        @test eltype((@with df @replace x = 1 @if _n == 1).x) <: AbstractFloat
+        df = DataFrame(x=[1, 2, missing])
+        @test eltype((@with df @replace x = 1.1 @if _n == 1).x) <: Union{T, Missing} where T <: AbstractFloat
+    end
+
     @testset "Error handling" begin
         @test_throws Exception @with df @replace y = 1
     end