add an option to intersect arguments passed to Cols

NEWS.md

@@ -7,6 +7,8 @@
 * 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 `operation` keyword argument in `Cols`
+  ([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,20 @@ 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[]
+    if idx.operation == :union
+        return union(getindex.(Ref(x), idx.cols)...)
+    elseif idx.operation == :intersect
+        return intersect(getindex.(Ref(x), idx.cols)...)
+    else
+        throw(ArgumentError("Unsupported operation :$(idx.operation). Either " *
+                            ":union or :intersect is required"))
+    end
+end
+# the definition below is needed because `:` is 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,67 @@ 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(operation=:intersect))) == 0
+    @test ncol(df[:, Cols(operation=:intersect)]) == 0
+
+    @test select(df, Cols(:, operation=:intersect)) == df[:, :]
+    @test df[:, Cols(, operation=:intersect:)] == df[:, :]
+
+    @test select(df, Cols(1, 2, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(1, :b, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(:a, 2, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(:a, :b, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, 1, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(:b, 1, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, :a, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(:b, :a, operation=:intersect)) == df[:, 2:1]
+
+    @test df[:, Cols(1, 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(1, :b, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, :b, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:b, 1, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, :a, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:b, :a, operation=:intersect)] == df[:, 2:1]
+
+    @test df[:, Cols(1, 1, 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 1, :b, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, 2, :b, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(:a, :b, 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, :a, operation=:intersect)] == df[:, 2:1]
+
+    @test select(df, Cols(1, "b", operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols("a", 2, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols("a", "b", operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols("b", 1, operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols(2, "a", operation=:intersect)) == df[:, 2:1]
+    @test select(df, Cols("b", "a", operation=:intersect)) == df[:, 2:1]
+
+    @test df[:, Cols(1, "b", operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", "b", operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("b", 1, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, "a", operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("b", "a", operation=:intersect)] == df[:, 2:1]
+
+    @test df[:, Cols("a", 1, "b", operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", 2, "b", operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols("a", "b", 2, operation=:intersect)] == df[:, 2:1]
+    @test df[:, Cols(2, 1, "a", operation=:intersect)] == df[:, 2:1]
+
+    df = DataFrame(a1=1, a2=2, b1=3, b2=4)
+    @test df[:, Cols(r"a", Not(r"1"), operation=:intersect)] == df[:, [:a2]]
+    @test df[:, Cols(Not(r"1"), r"a", operation=:intersect)] == df[:, [:a2]]
+    @test df[:, Cols(4:-1:1, 1:3, operation=:intersect)] == df[:, [3, 2]]
+    @test df[:, Cols(1:3, 4:-1:1, 1:3, operation=:intersect)] == df[:, [2, 3]]
+    @test df[:, Cols(1:3, 4:-1:1, 1:3, :b1, operation=:intersect)] == df[:, [:b1]]
+    @test df[:, Cols(x -> x[1] == 'a', operation=:intersect)] == df[:, [1, 2]]
+    @test df[:, Cols(x -> x[end] == '1', operation=:intersect)] == df[:, [1, 3]]
+    @test df[:, Cols(x -> x[end] == '3', operation=:intersect)] == DataFrame()
+    @test_throws ArgumentError df[:, Cols(x -> true, 1, operation=:intersect)]
+    @test_throws ArgumentError df[:, Cols(1, x -> true, operation=:intersect)]
 end
 
 @testset "views" begin