Support reading from and writing to Arrow files (#415)

This requires overriding `Arrow.DictEncoding` so that an `Arrow.DictEncoded`
with a `CategoricalArray` dictionary with one entry per level is created.
This is the only way to ensure that indexing the Arrow column gives
`CategoricalValue` objects. In practice such columns will most often
be used after conversion to `CategoricalArray` via `copy`, `DataFrame`, etc.

Apparently, pandas do not allow reading the resulting file if the array
allows for missing values as it does not accept `missing` in the dictionary.
Instead it would need missing entries to be coded via null indices, which
is less efficient.

Require Julia 1.6 as tests fail on older Julia versions.

Author: nalimilan

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
+          - '1.6'
           - '1' # automatically expands to the latest stable 1.x release of Julia
           - 'nightly'
         os:

Project.toml

@@ -13,18 +13,21 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [weakdeps]
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 SentinelArrays = "91c51154-3ec4-41a3-a24f-3f23e20d615c"
 StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 
 [extensions]
+CategoricalArraysArrowExt = "Arrow"
 CategoricalArraysJSONExt = "JSON"
 CategoricalArraysRecipesBaseExt = "RecipesBase"
 CategoricalArraysSentinelArraysExt = "SentinelArrays"
 CategoricalArraysStructTypesExt = "StructTypes"
 
 [compat]
+Arrow = "2"
 Compat = "3.37, 4"
 DataAPI = "1.6"
 JSON = "0.15, 0.16, 0.17, 0.18, 0.19, 0.20, 0.21"
@@ -35,9 +38,10 @@ Requires = "1"
 SentinelArrays = "1"
 Statistics = "1"
 StructTypes = "1"
-julia = "1"
+julia = "1.6"
 
 [extras]
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
@@ -49,4 +53,4 @@ StructTypes = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Dates", "JSON", "JSON3", "Plots", "PooledArrays", "RecipesBase", "SentinelArrays", "StructTypes", "Test"]
+test = ["Arrow", "Dates", "JSON", "JSON3", "Plots", "PooledArrays", "RecipesBase", "SentinelArrays", "StructTypes", "Test"]

ext/CategoricalArraysArrowExt.jl

@@ -0,0 +1,70 @@
+module CategoricalArraysArrowExt
+
+using CategoricalArrays
+import Arrow
+import Arrow: ArrowTypes
+
+const CATARRAY_ARROWNAME = Symbol("JuliaLang.CategoricalArrays.CategoricalArray")
+ArrowTypes.arrowname(::Type{<:CategoricalValue}) = CATARRAY_ARROWNAME
+ArrowTypes.arrowmetadata(::Type{CategoricalValue{T, R}}) where {T, R} = string(R)
+
+ArrowTypes.arrowname(::Type{Union{<:CategoricalValue, Missing}}) = CATARRAY_ARROWNAME
+ArrowTypes.arrowmetadata(::Type{Union{CategoricalValue{T, R}, Missing}}) where {T, R} =
+    string(R)
+
+const REFTYPES = Dict(string(T) => T for T in (Int128, Int16, Int32, Int64, Int8, UInt128,
+                                               UInt16, UInt32, UInt64, UInt8))
+function ArrowTypes.JuliaType(::Val{CATARRAY_ARROWNAME},
+                              ::Type{S}, meta::String) where S
+    R = REFTYPES[meta]
+    return CategoricalValue{S, R}
+end
+
+for (MV, MT) in ((:V, :T), (:(Union{V,Missing}), :(Union{T,Missing})))
+    @eval begin
+        function Arrow.DictEncoding{$MV,S,A}(id, data::Arrow.List{U, O, B},
+                                             isOrdered, metadata) where
+            {T, R, V<:CategoricalValue{T,R}, S, O, A, B, U}
+            newdata = Arrow.List{$MT,O,B}(data.arrow, data.validity, data.offsets,
+                                          data.data, data.ℓ, data.metadata)
+            levels = Missing <: $MT ? collect(skipmissing(newdata)) : newdata
+            catdata = CategoricalVector{$MT,R}(newdata, levels=levels)
+            return Arrow.DictEncoding{$MV,S,typeof(catdata)}(id, catdata,
+                                                             isOrdered, metadata)
+        end
+
+        function Arrow.DictEncoding{$MV,S,A}(id, data::Arrow.Primitive{U, B},
+                                            isOrdered, metadata) where
+            {T, R, V<:CategoricalValue{T,R}, S, A, B, U}
+            newdata = Arrow.Primitive{$MT,B}(data.arrow, data.validity, data.data,
+                                            data.ℓ, data.metadata)
+            levels = Missing <: $MT ? collect(skipmissing(newdata)) : newdata
+            catdata = CategoricalVector{$MT,R}(newdata, levels=levels)
+            return Arrow.DictEncoding{$MV,S,typeof(catdata)}(id, catdata,
+                                                             isOrdered, metadata)
+        end
+    end
+end
+
+function Base.copy(x::Arrow.DictEncoded{V}) where {T, R, V<:CategoricalValue{T, R}}
+    pool = CategoricalPool{T,R}(x.encoding.data)
+    inds = x.indices
+    refs = similar(inds, R)
+    refs .= inds .+ one(R)
+    return CategoricalVector{T}(refs, pool)
+end
+
+function Base.copy(x::Arrow.DictEncoded{Union{Missing,V}}) where
+    {T, R, V<:CategoricalValue{T, R}}
+    ismissing(x.encoding.data[1]) ||
+        throw(ErrorException("`missing` must be the first value in a " *
+                             "`CategoricalArray` pool"))
+    levels = collect(skipmissing(x.encoding.data))
+    pool = CategoricalPool{T,R}(levels)
+    inds = x.indices
+    refs = similar(inds, R)
+    refs .= inds
+    return CategoricalVector{Union{T,Missing}}(refs, pool)
+end
+
+end

src/CategoricalArrays.jl

@@ -41,6 +41,7 @@ module CategoricalArrays
 
     @static if !isdefined(Base, :get_extension)
         function __init__()
+            @require Arrow="69666777-d1a9-59fb-9406-91d4454c9d45" include("../ext/CategoricalArraysArrowExt.jl")
             @require JSON="682c06a0-de6a-54ab-a142-c8b1cf79cde6" include("../ext/CategoricalArraysJSONExt.jl")
             @require RecipesBase="3cdcf5f2-1ef4-517c-9805-6587b60abb01" include("../ext/CategoricalArraysRecipesBaseExt.jl")
             @require SentinelArrays="91c51154-3ec4-41a3-a24f-3f23e20d615c" include("../ext/CategoricalArraysSentinelArraysExt.jl")

test/13_arraycommon.jl

@@ -10,6 +10,8 @@ using StructTypes
 using RecipesBase
 using Plots
 using SentinelArrays
+using Arrow
+using Missings
 
 const ≅ = isequal
 const ≇ = !isequal
@@ -2071,6 +2073,57 @@ StructTypes.StructType(::Type{<:MyCustomType}) = StructTypes.Struct()
     @test levels(readx.var) == levels(x.var)
 end
 
+if Int == Int64
+    @testset "writing and reading Arrow files" for f in (identity, passmissing(string))
+        xref = f.([3, 1, 4, 1, 4])
+        x = categorical(f.([3, 1, 4, 1, 4]))
+        tbl = mktemp() do path, io
+            Arrow.write(path, (x=x,))
+            Arrow.Table(path)
+        end
+        @test tbl.x == x
+        @test tbl.x isa Arrow.DictEncoded{CategoricalValue{eltype(xref), UInt32}, Int8,
+                                          <: CategoricalVector{eltype(xref), UInt32}}
+        @test copy(tbl.x) == x
+        @test copy(x) isa CategoricalArray{eltype(xref),1,UInt32}
+
+        x = categorical(f.([3, 1, 4, 1, 4]), compress=true)
+        tbl = mktemp() do path, io
+            Arrow.write(path, (x=x,))
+            Arrow.Table(path)
+        end
+        @test tbl.x == x
+        @test tbl.x isa Arrow.DictEncoded{CategoricalValue{eltype(xref), UInt8}, Int8,
+                                          <: CategoricalVector{eltype(xref), UInt8}}
+        @test copy(tbl.x) == x
+        @test copy(x) isa CategoricalArray{eltype(xref),1,UInt8}
+
+        x = categorical(recode(xref, 1 => missing))
+        tbl = mktemp() do path, io
+            Arrow.write(path, (x=x,))
+            Arrow.Table(path)
+        end
+        @test tbl.x ≅ x
+        @test tbl.x isa Arrow.DictEncoded{Union{CategoricalValue{eltype(xref), UInt32}, Missing},
+                                          Int8,
+                                          <: CategoricalVector{Union{eltype(xref), Missing},
+                                                               UInt32}}
+        @test copy(tbl.x) ≅ x
+        @test copy(x) isa CategoricalArray{Union{eltype(xref), Missing},1,UInt32}
+
+        recode!(x, missing => f(1))
+        tbl = mktemp() do path, io
+            Arrow.write(path, (x=x,))
+            Arrow.Table(path)
+        end
+        @test tbl.x == x
+        @test tbl.x isa Arrow.DictEncoded{Union{CategoricalValue{eltype(xref), UInt32}, Missing}, Int8,
+                                          <: CategoricalVector{Union{eltype(xref), Missing}, UInt32}}
+        @test copy(tbl.x) == x
+        @test copy(x) isa CategoricalArray{Union{eltype(xref), Missing},1,UInt32}
+    end
+end
+
 @testset "refarray, refvalue, refpool, and invrefpool" begin
     for y in (categorical(["b", "a", "c", "b"]),
               view(categorical(["a", "a", "c", "b"]), 1:3),