Return example data as VecColumnTable with Vector columns (#29)

Author: junyuan-chen

src/DiffinDiffsBase.jl

@@ -133,9 +133,9 @@ export cb,
        lincom,
        rescale
 
+include("tables.jl")
 include("utils.jl")
 include("time.jl")
-include("tables.jl")
 include("ScaledArrays.jl")
 include("treatments.jl")
 include("parallels.jl")

src/StatsProcedures.jl

@@ -436,7 +436,7 @@ function show(io::IO, ::MIME"text/plain", sp::StatsSpec{T}) where T
     _show_args(io, sp)
 end
 
-function _count!(objcount::IdDict, obj)
+function _count!(objcount::AbstractDict, obj)
     count = get(objcount, obj, 0)
     objcount[obj] = count + 1
 end

src/operations.jl

@@ -11,7 +11,7 @@ end
 function _refs_pool(col::AbstractArray, reftype::Type{<:Integer}=UInt32)
     refs = refarray(col)
     pool = refpool(col)
-    labeled = pool !== nothing
+    labeled = pool !== nothing && !(pool isa RotatingTimeRange)
     if !labeled
         refs, invpool, pool = _label(col, eltype(col), reftype)
     end

src/utils.jl

@@ -144,7 +144,13 @@ function exampledata(name::Union{Symbol,String})
     Symbol(name) in exampledata() ||
         throw(ArgumentError("example dataset $(name) does not exist"))
     path = (@__DIR__)*"/../data/$(name).csv.gz"
-    return open(path) |> GzipDecompressorStream |> read |> CSV.File
+    cols = open(path) |> GzipDecompressorStream |> read |> CSV.File |> VecColumnTable
+    # Avoid customized index type for columns from CSV.File
+    for i in 1:ncol(cols)
+        col = _columns(cols)[i]
+        _columns(cols)[i] = convert(Vector, col)
+    end
+    return cols
 end
 
 # Check whether the input data is a column table

test/runtests.jl

@@ -23,9 +23,9 @@ import DiffinDiffsBase: required, valid_didargs, result
 include("testutils.jl")
 
 const tests = [
+    "tables",
     "utils",
     "time",
-    "tables",
     "ScaledArrays",
     "terms",
     "treatments",

test/tables.jl

@@ -1,5 +1,5 @@
 @testset "VecColumnTable" begin
-    hrs = exampledata("hrs")
+    hrs = DataFrame(exampledata("hrs"))
     cols = VecColumnTable(AbstractVector[], Symbol[], Dict{Symbol,Int}())
     cols1 = VecColumnTable(AbstractVector[], Symbol[])
     @test cols1 == cols

test/utils.jl

@@ -23,9 +23,11 @@ end
 
 @testset "exampledata" begin
     @test exampledata() == (:hrs, :nsw, :mpdta)
-    @test size(exampledata(:hrs)) == (3280,)
-    @test size(exampledata(:nsw)) == (32834,)
-    @test size(exampledata(:mpdta)) == (2500,)
+    hrs = exampledata("hrs")
+    @test size(hrs) == (3280, 11)
+    @test hrs[1] isa Vector
+    @test size(exampledata(:nsw)) == (32834, 11)
+    @test size(exampledata(:mpdta)) == (2500, 5)
 end
 
 @testset "checktable" begin