Allow predicate in Cols

NEWS.md

@@ -70,6 +70,8 @@
 * the `DataFrame` constructor when matrix is passed to it as a first
   argument now allows `copycols` keyword argument
   ([#2829](https://github.com/JuliaData/DataFrames.jl/pull/2859))
+* `Cols` now accepts a predicate accepting column names as strings.
+  ([#2880](https://github.com/JuliaData/DataFrames.jl/pull/2880))
 
 ## Bug fixes
 

docs/src/lib/internals.md

@@ -16,4 +16,5 @@ getmaxwidths
 ourshow
 ourstrwidth
 @spawn_for_chunks
+Cols
 ```

src/other/index.jl

@@ -221,6 +221,17 @@ end
     isempty(idx.cols) ? (1:length(x)) : throw(ArgumentError("All(args...) is not supported: use Cols(args...) instead"))
 @inline Base.getindex(x::AbstractIndex, idx::Cols) =
     isempty(idx.cols) ? Int[] : union(getindex.(Ref(x), idx.cols)...)
+@inline Base.getindex(x::AbstractIndex, idx::Cols{Tuple{typeof(:)}}) = x[:]
+@inline Base.getindex(x::AbstractIndex, idx::Cols{<:Tuple{Function}}) =
+    findall(idx.cols[1], names(x))
+
+"""
+    Cols(f::Function)
+
+Select the columns whose names passed to the `f` predicate as strings return `true`.
+As a special case if `:` is passed (the `Colon` function) select all columns.
+"""
+Cols
 
 @inline function Base.getindex(x::AbstractIndex, idx::AbstractVector{<:Integer})
     if any(v -> v isa Bool, idx)

test/index.jl

@@ -474,6 +474,11 @@ end
     df = DataFrame(a1=1, a2=2, b1=3, b2=4)
     @test df[:, Cols(r"a", Not(r"1"))] == df[:, [1, 2, 4]]
     @test df[:, Cols(Not(r"1"), r"a")] == df[:, [2, 4, 1]]
+    @test df[:, Cols(x -> x[1] == 'a')] == df[:, [1, 2]]
+    @test df[:, Cols(x -> x[end] == '1')] == df[:, [1, 3]]
+    @test df[:, Cols(x -> x[end] == '3')] == DataFrame()
+    @test_throws MethodError df[:, Cols(x -> true, 1)] == DataFrame()
+    @test_throws MethodError df[:, Cols(1, x -> true)] == DataFrame()
 end
 
 @testset "views" begin