add an option to intersect arguments passed to Cols (#3224)

NEWS.md

@@ -7,6 +7,9 @@
 * Add `allunique` and allow transformations in `cols` argument of `describe`
   and `nonunique` when working with `SubDataFrame`
   ([3232](https://github.com/JuliaData/DataFrames.jl/pull/3232))
+* Add support for `operator` keyword argument in `Cols`
+  to take a set operation to apply to passed selectors (`union` by default)
+  ([3224](https://github.com/JuliaData/DataFrames.jl/pull/3224))
 
 # DataFrames.jl v1.4.4 Patch Release Notes
 

Project.toml

@@ -27,7 +27,7 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 [compat]
 CategoricalArrays = "0.10.0"
 Compat = "4.2"
-DataAPI = "1.13.0"
+DataAPI = "1.14.0"
 InlineStrings = "1.3.0"
 InvertedIndices = "1"
 IteratorInterfaceExtensions = "0.1.1, 1"

src/other/index.jl

@@ -229,8 +229,13 @@ end
 @inline Base.getindex(x::AbstractIndex, idx::Between) = x[idx.first]:x[idx.last]
 @inline Base.getindex(x::AbstractIndex, idx::All) =
     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 function Base.getindex(x::AbstractIndex, idx::Cols)
+    isempty(idx.cols) && return Int[]
+    return idx.operator(getindex.(Ref(x), idx.cols)...)
+end
+
+# the definition below is needed because `:` is a Function
 @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))

test/index.jl

@@ -479,6 +479,68 @@ end
     @test df[:, Cols(x -> x[end] == '3')] == DataFrame()
     @test_throws ArgumentError df[:, Cols(x -> true, 1)]
     @test_throws ArgumentError df[:, Cols(1, x -> true)]
+
+    @test ncol(select(df, Cols(operator=intersect))) == 0
+    @test ncol(df[:, Cols(operator=intersect)]) == 0
+
+    @test select(df, Cols(:, operator=intersect)) == df[:, :]
+    @test df[:, Cols(:, operator=intersect)] == df[:, :]
+
+    df = DataFrame(a=1, b=2, c=3)
+    @test select(df, Cols(1, 2, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(1, :b, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(:a, 2, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(:a, :b, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, 1, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(:b, 1, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, :a, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(:b, :a, operator=intersect)) == df[:, 2:1]
+
+    @test df[:, Cols(1, 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(1, :b, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, :b, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:b, 1, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, :a, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:b, :a, operator=intersect)] == df[:, 2:1]
+
+    @test df[:, Cols(1, 1, 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 1, :b, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 2, :b, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, :b, 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, :a, operator=intersect)] == df[:, 2:1]
+
+    @test select(df, Cols(1, "b", operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols("a", 2, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols("a", "b", operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols("b", 1, operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, "a", operator=intersect)) == df[:, 2:1]
+    @test select(df, Cols("b", "a", operator=intersect)) == df[:, 2:1]
+
+    @test df[:, Cols(1, "b", operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", "b", operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("b", 1, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, "a", operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("b", "a", operator=intersect)] == df[:, 2:1]
+
+    @test df[:, Cols("a", 1, "b", operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", 2, "b", operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", "b", 2, operator=intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, "a", operator=intersect)] == df[:, 2:1]
+
+    df = DataFrame(a1=1, a2=2, b1=3, b2=4)
+    @test df[:, Cols(r"a", Not(r"1"), operator=intersect)] == df[:, [:a2]]
+    @test df[:, Cols(Not(r"1"), r"a", operator=intersect)] == df[:, [:a2]]
+    @test df[:, Cols(4:-1:1, 2:3, operator=intersect)] == df[:, [3, 2]]
+    @test df[:, Cols(1:3, 4:-1:1, 2:3, operator=intersect)] == df[:, [2, 3]]
+    @test df[:, Cols(1:3, 4:-1:1, 1:3, :b1, operator=intersect)] == df[:, [:b1]]
+    @test df[:, Cols(x -> x[1] == 'a', operator=intersect)] == df[:, [1, 2]]
+    @test df[:, Cols(x -> x[end] == '1', operator=intersect)] == df[:, [1, 3]]
+    @test df[:, Cols(x -> x[end] == '3', operator=intersect)] == DataFrame()
+    @test_throws ArgumentError df[:, Cols(x -> true, 1, operator=intersect)]
+    @test_throws ArgumentError df[:, Cols(1, x -> true, operator=intersect)]
 end
 
 @testset "views" begin