Add extension for Arrow + compatibility for JSON.Object handling

Author: quinnj

Project.toml

@@ -3,6 +3,7 @@ uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 version = "1.0.0"
 
 [deps]
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Mmap = "a63ad114-7e13-5084-954f-fe012c677804"
@@ -12,15 +13,25 @@ StructUtils = "ec057cc2-7a8d-4b58-b3b3-92acb9f63b42"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
+[weakdeps]
+ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd"
+
+[extensions]
+JSONArrowExt = ["ArrowTypes"]
+
 [compat]
+Arrow = "2.8.0"
+ArrowTypes = "2.2"
 Parsers = "1, 2"
 PrecompileTools = "1"
 StructUtils = "2"
 julia = "1.9"
 
 [extras]
+Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
+ArrowTypes = "31f734f8-188a-4ce0-8406-c8a06bd891cd"
 Tar = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Tar", "Test"]
+test = ["Arrow", "Tar", "Test"]

ext/JSONArrowExt.jl

@@ -0,0 +1,78 @@
+module JSONArrowExt
+
+using JSON, ArrowTypes
+
+const JSON_ARROW_NAME = Symbol("JuliaLang.JSON.Object")
+
+ArrowTypes.ArrowKind(::Type{<:JSON.Object}) = ArrowTypes.StructKind()
+
+ArrowTypes.toarrow(x::JSON.Object) = (; x...)
+ArrowTypes.arrowname(::Type{<:JSON.Object}) = JSON_ARROW_NAME
+ArrowTypes.JuliaType(::Val{JSON_ARROW_NAME}) = JSON.Object{Symbol,Any}
+
+ArrowTypes.fromarrowstruct(::Type{T}, ::Val{nms}, vals...) where {T <: JSON.Object, nms} =
+    T(nms[i] => vals[i] for i in 1:length(nms))
+
+ArrowTypes.ToArrow(x::AbstractArray{T}) where {T <: JSON.Object} = _toarrow(x)
+ArrowTypes.ToArrow(x::AbstractArray{Union{T,Missing}}) where {T <: JSON.Object} = _toarrow(x)
+ArrowTypes.ToArrow(x::JSON.Object) = toarrow(x)
+
+function _toarrow(x::Union{AbstractArray{T}, AbstractArray{Union{T,Missing}}}) where {T<:JSON.Object}
+    isempty(x) && return Missing[]
+    x isa AbstractArray{Missing} && return x
+    fields = JSON.Object{Symbol, Type}()
+    seen_fields = Set{Symbol}()
+    for (i, y) in enumerate(x)
+        y === missing && continue
+        current_fields = Set{Symbol}()
+        for (k, vv) in y
+            key = Symbol(k)
+            push!(current_fields, key)
+            v = toarrow(vv)
+            vtype = typeof(v)
+            existing_type = get(fields, key, nothing)
+            if existing_type !== nothing
+                if !(vtype <: existing_type)
+                    fields[key] = ArrowTypes.promoteunion(existing_type, vtype)
+                end
+            else
+                if i == 1
+                    fields[key] = vtype
+                else
+                    fields[key] = Union{vtype, Missing}
+                end
+            end
+        end
+        for field in seen_fields
+            if !(field in current_fields)
+                existing_type = fields[field]
+                if !(Missing <: existing_type)
+                    fields[field] = Union{existing_type, Missing}
+                end
+            end
+        end
+        union!(seen_fields, current_fields)
+    end
+    for y in x
+        obj = toarrow(y)
+        obj === missing && continue
+        for (field, _) in fields
+            if !haskey(obj, field)
+                existing_type = fields[field]
+                if !(Missing <: existing_type)
+                    fields[field] = Union{existing_type, Missing}
+                end
+            end
+        end
+    end
+    nms = Tuple(keys(fields))
+    NT = NamedTuple{nms, Tuple{values(fields)...}}
+    return ArrowTypes.ToArrow{NT,typeof(x)}(x)
+end
+
+function ArrowTypes._convert(::Type{NamedTuple{nms,T}}, nt) where {nms,T}
+    vals = Tuple((nt !== missing && haskey(nt, nm)) ? toarrow(getproperty(nt, nm)) : missing for nm in nms)
+    return NamedTuple{nms,T}(vals)
+end
+
+end # module

test/arrow.jl

@@ -0,0 +1,68 @@
+using Test, JSON, Arrow
+
+obj1 = JSON.parse("""
+{
+    "int": 1,
+    "float": 2.1
+}
+""")
+
+obj2 = JSON.parse("""
+{
+    "int": 1,
+    "float": 2.1,
+    "bool1": true,
+    "bool2": false,
+    "none": null,
+    "str": "\\"hey there sailor\\"",
+    "arr": [null, 1, "hey"],
+    "arr2": [1.2, 3.4, 5.6]
+}
+""")
+
+obj3 = JSON.parse("""
+{
+    "int": 1,
+    "float": 2.1,
+    "bool1": true,
+    "bool2": false,
+    "none": null,
+    "str": "\\"hey there sailor\\"",
+    "obj": {
+                "a": 1,
+                "b": null,
+                "c": [null, 1, "hey"],
+                "d": [1.2, 3.4, 5.6]
+            },
+    "arr": [null, 1, "hey"],
+    "arr2": [1.2, 3.4, 5.6]
+}
+""")
+
+tbl = (; json=[obj1, obj2, obj3])
+
+arrow = Arrow.Table(Arrow.tobuffer(tbl))
+@test arrow.json[1].int == 1
+@test arrow.json[1].float == 2.1
+
+@test arrow.json[2].int == 1
+@test arrow.json[2].float == 2.1
+@test arrow.json[2].bool1 == true
+@test arrow.json[2].bool2 == false
+@test arrow.json[2].none === missing
+@test arrow.json[2].str == "\"hey there sailor\""
+@test isequal(arrow.json[2].arr, [missing, 1, "hey"])
+@test arrow.json[2].arr2 == [1.2, 3.4, 5.6]
+
+@test arrow.json[3].int == 1
+@test arrow.json[3].float == 2.1
+@test arrow.json[3].bool1 == true
+@test arrow.json[3].bool2 == false
+@test arrow.json[3].none === missing
+@test arrow.json[3].str == "\"hey there sailor\""
+@test arrow.json[3].obj.a == 1
+@test arrow.json[3].obj.b === nothing
+@test isequal(arrow.json[3].obj.c, [missing, 1, "hey"])
+@test arrow.json[3].obj.d == [1.2, 3.4, 5.6]
+@test isequal(arrow.json[3].arr, [missing, 1, "hey"])
+@test arrow.json[3].arr2 == [1.2, 3.4, 5.6]

test/runtests.jl

@@ -4,6 +4,7 @@ include(joinpath(dirname(pathof(JSON)), "../test/object.jl"))
 include(joinpath(dirname(pathof(JSON)), "../test/lazy.jl"))
 include(joinpath(dirname(pathof(JSON)), "../test/parse.jl"))
 include(joinpath(dirname(pathof(JSON)), "../test/json.jl"))
+include(joinpath(dirname(pathof(JSON)), "../test/arrow.jl"))
 
 function tar_files(tarball::String)
     data = Dict{String, Vector{UInt8}}()