Improve performance of dropmissing

src/abstractdataframe/abstractdataframe.jl

@@ -970,6 +970,7 @@ julia> dropmissing(df, [:x, :y])
 @inline function dropmissing(df::AbstractDataFrame,
                              cols::Union{ColumnIndex, MultiColumnIndex}=:;
                              view::Bool=false, disallowmissing::Bool=!view)
+    # Changes to be made here
     rowidxs = completecases(df, cols)
     if view
         if disallowmissing
@@ -983,6 +984,36 @@ julia> dropmissing(df, [:x, :y])
     end
 end
 
+@inline function _getindex_disallowmissing(df::AbstractDataFrame,
+    row_inds::AbstractVector{T},
+    col_disallowmissing_inds::Union{ColumnIndex,MultiColumnIndex}) where {T<:Integer}
+
+    newdf = df[row_inds, :]
+    disallowmissing!(newdf, col_disallowmissing_inds)
+
+    return df
+end
+
+# new implementation; to be swapped with the above if design is okay
+@inline function _dropmissing(df::AbstractDataFrame,
+    cols::Union{ColumnIndex,MultiColumnIndex}=:;
+    view::Bool=false, disallowmissing::Bool=!view)
+    rowidxs = completecases(df, cols)
+    if view
+        if disallowmissing
+            throw(ArgumentError("disallowmissing=true is incompatible with view=true"))
+        end
+        return Base.view(df, rowidxs, :)
+    else
+        if disallowmissing
+            newdf = _getindex_disallowmissing(df, rowidxs, cols)
+        else
+            newdf = df[rowidxs, :]
+        end
+        return newdf
+    end
+end
+
 """
     dropmissing!(df::AbstractDataFrame, cols=:; disallowmissing::Bool=true)
 

src/dataframe/dataframe.jl

@@ -625,6 +625,45 @@ Base.getindex(df::DataFrame, row_ind::Colon, col_inds::MultiColumnIndex) =
 Base.getindex(df::DataFrame, row_ind::typeof(!), col_inds::MultiColumnIndex) =
     select(df, index(df)[col_inds], copycols=false)
 
+# Gets a subset of a DataFrame by `row_inds` and disallows missing values in columns specified by `col_disallowmissing_inds` 
+# Called by dropmissing()
+function _getindex_disallowmissing(df::DataFrame, row_inds::AbstractVector{T},
+    col_disallowmissing_inds::Union{ColumnIndex,MultiColumnIndex,Colon}) where {T<:Integer}
+
+    # With knowledge of exact indices we can skip a lot of loop iterations (as opposed to filtering by a mask)
+    selected_rows = T === Bool ? _findall(row_inds) : row_inds
+    new_columns = Vector{AbstractVector}(undef, ncol(df))
+    df_columns = DataFrames._columns(df)
+
+    # We need to know the exact indices of columns that disallowmissing should be applied to
+    disallowmissing_columns_inds = index(df)[col_disallowmissing_inds]
+
+    # Threading decision rule borrowed from `_threaded_getindex`
+    if Threads.nthreads() > 1 && ncol(df) > 1 && length(selected_rows) >= 1_000_000
+        @sync for i in eachindex(new_columns)
+            # for each column, check if disallowmissing should be applied
+            Threads.@spawn if i in disallowmissing_columns_inds
+                @inbounds new_columns[i] = _getindex_disallowmissing(df_columns[i], selected_rows)
+            else
+                @inbounds new_columns[i] = df_columns[i][selected_rows]
+            end
+        end
+    else
+        @inbounds for i in eachindex(new_columns)
+            if i in disallowmissing_columns_inds
+                new_columns[i] = _getindex_disallowmissing(df_columns[i], selected_rows)
+            else
+                new_columns[i] = df_columns[i][selected_rows]
+            end
+        end
+    end
+
+    newdf = DataFrame(new_columns, copy(DataFrames.index(df)), copycols=false)
+    _copy_all_note_metadata!(newdf, df)
+
+    return newdf
+end
+
 ##############################################################################
 ##
 ## setindex!()

src/other/utils.jl

@@ -477,3 +477,22 @@ function _findall(B::BitVector)::Union{UnitRange{Int}, Vector{Int}}
     end
     @assert false "should not be reached"
 end
+
+# Creates a vector like `vect` without the missing type and copies only data at indices in `inds`
+# Fuses two operations that were previously done by getindex + disallowmissing
+# unsafe because it does not check for `inds` values to be not missing (ensured by the caller)
+# Called by `_getindex_disallowmissing(df,...)`
+function _getindex_disallowmissing(vect::AbstractVector{T}, inds::AbstractVector{<:Integer}) where {T}
+    # General fallback, relies on custom vectors' own methods
+    return vect[inds]|>disallowmissing
+end
+
+function _getindex_disallowmissing(vect::Vector{T}, inds::AbstractVector{<:Integer}) where {T}
+    output = Vector{nonmissingtype(T)}(undef, length(inds))
+    last_ind = 0
+    @inbounds for i in eachindex(inds)
+        last_ind += 1
+        output[last_ind] = vect[inds[i]]
+    end
+    return output
+end