Some optimizations in the 'DropMissing' transform

Author: eliascarv

src/transforms/dropmissing.jl

@@ -43,24 +43,19 @@ DropMissing(cols::T...) where {T<:Col} = DropMissing(colspec(cols))
 
 isrevertible(::Type{<:DropMissing}) = true
 
-_ftrans(::DropMissing{AllSpec}, snames) = Filter(row -> all(!ismissing, row))
-_ftrans(::DropMissing, snames) = Filter(row -> all(!ismissing, row[nm] for nm in snames))
-
-# nonmissing 
-_nonmissing(::Type{T}, x) where {T} = x
-_nonmissing(::Type{Union{Missing,T}}, x) where {T} = collect(T, x)
-_nonmissing(::Type{Missing}, x) = []
-_nonmissing(x) = _nonmissing(eltype(x), x)
-
 function preprocess(transform::DropMissing, table)
-  schema = Tables.schema(table)
-  names = schema.names
+  names = Tables.schema(table).names
   snames = choose(transform.colspec, names)
-  ftrans = _ftrans(transform, snames)
+  ftrans = Filter(row -> all(!ismissing(row[nm]) for nm in snames))
   fprep = preprocess(ftrans, table)
   ftrans, fprep, snames
 end
 
+_nonmissing(x) = _nonmissing(eltype(x), x)
+_nonmissing(::Type{T}, x) where {T} = x
+_nonmissing(::Type{Missing}, x) = []
+_nonmissing(::Type{Union{Missing,T}}, x) where {T} = collect(T, x)
+
 function applyfeat(::DropMissing, feat, prep)
   # apply filter transform
   ftrans, fprep, snames = prep
@@ -69,28 +64,28 @@ function applyfeat(::DropMissing, feat, prep)
   # drop Missing type
   cols = Tables.columns(newfeat)
   names = Tables.columnnames(cols)
-  columns = map(names) do nm
-    x = Tables.getcolumn(cols, nm)
-    nm ∈ snames ? _nonmissing(x) : x
+  columns = map(names) do name
+    x = Tables.getcolumn(cols, name)
+    name ∈ snames ? _nonmissing(x) : x
   end
   𝒯 = (; zip(names, columns)...)
   newfeat = 𝒯 |> Tables.materializer(feat)
 
   # original column types
   types = Tables.schema(feat).types
 
-  newfeat, (ftrans, ffcache, types)
+  newfeat, (ftrans, ffcache, snames, types)
 end
 
 function revertfeat(::DropMissing, newfeat, fcache)
-  ftrans, ffcache, types = fcache
+  ftrans, ffcache, snames, types = fcache
 
   # reintroduce Missing type
   cols = Tables.columns(newfeat)
   names = Tables.columnnames(cols)
-  columns = map(zip(types, names)) do (T, nm)
-    x = Tables.getcolumn(cols, nm)
-    collect(T, x)
+  columns = map(types, names) do T, name
+    x = Tables.getcolumn(cols, name)
+    name ∈ snames ? collect(T, x) : x
   end
   𝒯 = (; zip(names, columns)...)
   ofeat = 𝒯 |> Tables.materializer(newfeat)