support of deleteat! for subdatasets

Author: sl-solution

src/dataset/del_and_append.jl

@@ -27,7 +27,7 @@ julia> deleteat!(ds, 2)
 ```
 
 """
-function Base.deleteat!(ds::Dataset, inds)
+function Base.deleteat!(ds::AbstractDataset, inds)
 
 # Modify Dataset
     if !isempty(inds) && size(ds, 2) == 0
@@ -40,7 +40,7 @@ function Base.deleteat!(ds::Dataset, inds)
 end
 
 # Modify Dataset
-function Base.deleteat!(ds::Dataset, inds::AbstractVector{Bool})
+function Base.deleteat!(ds::AbstractDataset, inds::AbstractVector{Bool})
     if length(inds) != size(ds, 1)
         throw(BoundsError(ds, (inds, :)))
     end
@@ -49,13 +49,16 @@ function Base.deleteat!(ds::Dataset, inds::AbstractVector{Bool})
 end
 
 # Modify Dataset
-Base.deleteat!(ds::Dataset, inds::Not) = deleteat!(ds, axes(ds, 1)[inds])
+Base.deleteat!(ds::AbstractDataset, inds::Not) = deleteat!(ds, axes(ds, 1)[inds])
 
 # Modify Dataset
-function _delete!_helper(ds::Dataset, drop)
+function _delete!_helper(ds::AbstractDataset, drop)
     cols = _columns(ds)
     isempty(cols) && return ds
-
+    if ds isa SubDataset
+        deleteat!(getfield(ds, :rows), drop)
+        return ds
+    end
     n = nrow(ds)
     col1 = cols[1]
     deleteat!(col1, drop)

src/subdataset/subdataset.jl

@@ -110,8 +110,14 @@ Base.@propagate_inbounds function SubDataset(parent::Dataset, rows::AbstractVect
     end
     SubDataset(parent, SubIndex(index(parent), cols), rows)
 end
+Base.@propagate_inbounds function SubDataset(parent::Dataset, rows::AbstractUnitRange, cols)
+    @boundscheck if !checkindex(Bool, axes(parent, 1), rows)
+        throw(BoundsError(parent, (rows, cols)))
+    end
+    SubDataset(parent, SubIndex(index(parent), cols), convert(Vector{Int}, rows))
+end
 Base.@propagate_inbounds SubDataset(parent::Dataset, ::Colon, cols) =
-    SubDataset(parent, axes(parent, 1), cols)
+    SubDataset(parent, collect(axes(parent, 1)), cols)
 @inline SubDataset(parent::Dataset, row::Integer, cols) =
     throw(ArgumentError("invalid row index: $row of type $(typeof(row))"))
 
@@ -275,8 +281,8 @@ Base.setproperty!(::SubDataset, ::AbstractString, ::Any) =
 
 Base.copy(sds::SubDataset) = parent(sds)[rows(sds), parentcols(index(sds), :)]
 
-Base.deleteat!(ds::SubDataset, ind) =
-    throw(ArgumentError("SubDataset does not support deleting rows"))
+# Base.deleteat!(ds::SubDataset, ind) =
+    # throw(ArgumentError("SubDataset does not support deleting rows"))
 
 function Dataset(sds::SubDataset)
 

test/data.jl

@@ -117,6 +117,13 @@ using Test, InMemoryDatasets, Random, CategoricalArrays
         ds2 = ds[nrow(ds):-1:1, ncol(ds):-1:1]
         @test unique(ds2, cols, keep = keepval, mapformats = mfmt) == unique(view(ds, nrow(ds):-1:1, ncol(ds):-1:1), cols, keep = keepval, mapformats = mfmt)
     end
+
+    ds1 = Dataset(a = Union{String, Missing}["a", "b", "a", "b", "a", "b"],
+                    b = Vector{Union{Int, Missing}}(1:6),
+                    c = Union{Int, Missing}[1:3;1:3])
+    ds = vcat(ds1, ds1)
+    unique!(ds, 1, keep = :random)
+    @test ds == unique(ds, 1)
 end
 
 @testset "completecases and dropmissing" begin

test/sort.jl

@@ -399,6 +399,10 @@ end
     sds = view(ds, [true, true, true, true, true, false, false, false, false, false], 1:2)
     sort!(sds, 1)
     @test sds == sort(sds, 1)
+
+    sds = view(ds, 1:5, 1:2)
+    sort!(sds, 1)
+    @test sds == sort(sds, 1)
 end
 
 @testset "issorted/issorted!" begin

test/subdataset.jl

@@ -222,12 +222,6 @@ end
     @test names(ds) == names(x)[[4, 2]]
 end
 
-@testset "deleteat!" begin
-    y = 1.0:10.0
-    ds = view(Dataset(y=y), 2:6, :)
-    @test_throws ArgumentError deleteat!(ds, 1)
-end
-
 @testset "parent" begin
     ds = Dataset(a=Union{Int, Missing}[1, 2, 3, 1, 2, 2],
                    b=[2.0, missing, 1.2, 2.0, missing, missing],
@@ -293,7 +287,29 @@ end
     if VERSION >= v"1.7.0-DEV"
         @test_throws ArgumentError SubDataset(sds, Integer[true, true, true], :)
     else
-        @test SubDataset(sds, Integer[true, true, true], :) ==
+        @test SubDataset(sds, Int[true, true, true], :) ==
               SubDataset(sds, [1, 1, 1], :)
     end
 end
+
+@testset "deletat! for views" begin
+    ds1 = Dataset(a = Union{String, Missing}["a", "b", "a", "b", "a", "b"],
+                    b = Vector{Union{Int, Missing}}(1:6),
+                    c = Union{Int, Missing}[1:3;1:3])
+    ds = vcat(ds1, ds1)
+    sds = view(ds, 1:4, [1,3])
+    deleteat!(sds, 1)
+    @test sds == ds[2:4, [1,3]]
+    sds = view(ds, 1:4, [1,3])
+    deleteat!(sds, 1:2)
+    @test sds == ds[3:4, [1,3]]
+    sds = view(ds, 1:4, [1,3])
+    deleteat!(sds, [1,3])
+    @test sds == ds[[2,4], [1,3]]
+    sds = view(ds, :, :)
+    deleteat!(sds, Not(1:3))
+    @test sds == ds[1:3, :]
+    sds = view(ds, :, :)
+    deleteat!(sds, [1,3,5,7,9,11])
+    @test sds == ds[[2,4,6,8,10,12], :]
+end