add findcols

NEWS.md

@@ -13,6 +13,9 @@
   columns only to a subset of the columns specified by the `cols`
   keyword argument
   ([#3386](https://github.com/JuliaData/DataFrames.jl/pull/3386))
+* add `findcols` that returns a vector of integer column indices
+  of a data frame that meet the passed condition function
+  ([#3389](https://github.com/JuliaData/DataFrames.jl/pull/3389))
 
 ## Bug fixes
 

docs/src/lib/functions.md

@@ -26,7 +26,7 @@ This is a list of operations that currently make use of multi-threading:
   * a transformation produces one row per group and the passed transformation
     is a custom function (i.e. not for standard reductions, which use
     optimized single-threaded methods).
-- `dropmissing` when the provided data frame has more than 1 column and `view=false` 
+- `dropmissing` when the provided data frame has more than 1 column and `view=false`
   (subsetting of individual columns is spawned in separate tasks).
 
 In general at least Julia 1.4 is required to ensure that multi-threading is used
@@ -170,6 +170,11 @@ unique
 unique!
 ```
 
+## Filtering columns
+```@docs
+findcols
+```
+
 ## Working with missing values
 ```@docs
 allowmissing

src/DataFrames.jl

@@ -67,6 +67,7 @@ export AbstractDataFrame,
        dropmissing!,
        dropmissing,
        fillcombinations,
+       findcols,
        flatten,
        groupby,
        groupindices,

src/abstractdataframe/abstractdataframe.jl

@@ -3252,3 +3252,29 @@ function Base.iterate(itr::Iterators.PartitionIterator{<:AbstractDataFrame}, sta
     r = min(state + itr.n - 1, last_idx)
     return view(itr.c, state:r, :), r + 1
 end
+
+"""
+    findall(f, df::AbstractDataFrame)
+
+Return an integer vector `I` of the column indices `i` of `df` where `f(df[:, i])` returns `true`.
+If there are no such columns of `df`, return `Int[]`.
+
+# Examples
+
+```jldoctest
+julia> df = DataFrame(a=[1, missing], b=[2, 3], c=[missing, 4])
+2×3 DataFrame
+ Row │ a        b      c
+     │ Int64?   Int64  Int64?
+─────┼─────────────────────────
+   1 │       1      2  missing
+   2 │ missing      3        4
+
+julia> findcols(x -> any(ismissing, x), df)
+2-element Vector{Int64}:
+ 1
+ 3
+```
+"""
+findcols(f::Function, df::AbstractDataFrame) =
+    findall(f, eachcol(df))

test/dataframe.jl

@@ -2387,4 +2387,17 @@ end
     @test eltype(collect(p)) <: DataFrames.DataFrameRows
 end
 
+@testset "findcols" begin
+    df = DataFrame(a=[1, missing], b=[2, 3], c=[missing, 4])
+    @test findcols(x -> any(ismissing, x), df) == [1, 3]
+    @test findcols(x -> true, df) == [1, 2, 3]
+    @test findcols(x -> false, df) == Int[]
+    @test_throws TypeError findcols(x -> 1, df)
+
+    @test findcols(x -> any(ismissing, x), view(df, :, [1, 2])) == [1]
+    @test findcols(x -> true, view(df, :, [1, 2])) == [1, 2]
+    @test findcols(x -> false, view(df, :, [1, 2])) == Int[]
+    @test_throws TypeError findcols(x -> 1, view(df, :, [1, 2]))
+end
+
 end # module