--- /dev/null
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module ArrowJSON
+
+using Mmap
+using StructTypes, JSON3, Tables, SentinelArrays, Arrow
+
+# read json files as "table"
+# write to arrow stream/file
+# read arrow stream/file back
+
+abstract type Type end
+Type() = Null("null")
+StructTypes.StructType(::Base.Type{Type}) = StructTypes.AbstractType()
+
+children(::Base.Type{T}) where {T} = Field[]
+
+mutable struct Int <: Type
+ name::String
+ bitWidth::Int64
+ isSigned::Base.Bool
+end
+
+Int() = Int("", 0, true)
+Type(::Base.Type{T}) where {T <: Integer} = Int("int", 8 * sizeof(T), T <: Signed)
+StructTypes.StructType(::Base.Type{Int}) = StructTypes.Mutable()
+function juliatype(f, x::Int)
+ T = x.bitWidth == 8 ? Int8 : x.bitWidth == 16 ? Int16 :
+ x.bitWidth == 32 ? Int32 : x.bitWidth == 64 ? Int64 : Int128
+ return x.isSigned ? T : unsigned(T)
+end
+
+struct FloatingPoint <: Type
+ name::String
+ precision::String
+end
+
+Type(::Base.Type{T}) where {T <: AbstractFloat} = FloatingPoint("floatingpoint", T == Float16 ? "HALF" : T == Float32 ? "SINGLE" : "DOUBLE")
+StructTypes.StructType(::Base.Type{FloatingPoint}) = StructTypes.Struct()
+juliatype(f, x::FloatingPoint) = x.precision == "HALF" ? Float16 : x.precision == "SINGLE" ? Float32 : Float64
+
+struct FixedSizeBinary <: Type
+ name::String
+ byteWidth::Int64
+end
+
+Type(::Base.Type{NTuple{N, UInt8}}) where {N} = FixedSizeBinary("fixedsizebinary", N)
+children(::Base.Type{NTuple{N, UInt8}}) where {N} = Field[]
+StructTypes.StructType(::Base.Type{FixedSizeBinary}) = StructTypes.Struct()
+juliatype(f, x::FixedSizeBinary) = NTuple{x.byteWidth, UInt8}
+
+struct Decimal <: Type
+ name::String
+ precision::Int32
+ scale::Int32
+end
+
+Type(::Base.Type{Arrow.Decimal{P, S, T}}) where {P, S, T} = Decimal("decimal", P, S)
+StructTypes.StructType(::Base.Type{Decimal}) = StructTypes.Struct()
+juliatype(f, x::Decimal) = Arrow.Decimal{x.precision, x.scale, Int128}
+
+mutable struct Timestamp <: Type
+ name::String
+ unit::String
+ timezone::Union{Nothing ,String}
+end
+
+Timestamp() = Timestamp("", "", nothing)
+unit(U) = U == Arrow.Meta.TimeUnit.SECOND ? "SECOND" :
+ U == Arrow.Meta.TimeUnit.MILLISECOND ? "MILLISECOND" :
+ U == Arrow.Meta.TimeUnit.MICROSECOND ? "MICROSECOND" : "NANOSECOND"
+Type(::Base.Type{Arrow.Timestamp{U, TZ}}) where {U, TZ} = Timestamp("timestamp", unit(U), TZ === nothing ? nothing : String(TZ))
+StructTypes.StructType(::Base.Type{Timestamp}) = StructTypes.Mutable()
+unitT(u) = u == "SECOND" ? Arrow.Meta.TimeUnit.SECOND :
+ u == "MILLISECOND" ? Arrow.Meta.TimeUnit.MILLISECOND :
+ u == "MICROSECOND" ? Arrow.Meta.TimeUnit.MICROSECOND : Arrow.Meta.TimeUnit.NANOSECOND
+juliatype(f, x::Timestamp) = Arrow.Timestamp{unitT(x.unit), x.timezone === nothing ? nothing : Symbol(x.timezone)}
+
+struct Duration <: Type
+ name::String
+ unit::String
+end
+
+Type(::Base.Type{Arrow.Duration{U}}) where {U} = Duration("duration", unit(U))
+StructTypes.StructType(::Base.Type{Duration}) = StructTypes.Struct()
+juliatype(f, x::Duration) = Arrow.Duration{unit%(x.unit)}
+
+struct Date <: Type
+ name::String
+ unit::String
+end
+
+Type(::Base.Type{Arrow.Date{U, T}}) where {U, T} = Date("date", U == Arrow.Meta.DateUnit.DAY ? "DAY" : "MILLISECOND")
+StructTypes.StructType(::Base.Type{Date}) = StructTypes.Struct()
+juliatype(f, x::Date) = Arrow.Date{x.unit == "DAY" ? Arrow.Meta.DateUnit.DAY : Arrow.Meta.DateUnit.MILLISECOND, x.unit == "DAY" ? Int32 : Int64}
+
+struct Time <: Type
+ name::String
+ unit::String
+ bitWidth::Int64
+end
+
+Type(::Base.Type{Arrow.Time{U, T}}) where {U, T} = Time("time", unit(U), 8 * sizeof(T))
+StructTypes.StructType(::Base.Type{Time}) = StructTypes.Struct()
+juliatype(f, x::Time) = Arrow.Time{unitT(x.unit), x.unit == "SECOND" || x.unit == "MILLISECOND" ? Int32 : Int64}
+
+struct Interval <: Type
+ name::String
+ unit::String
+end
+
+Type(::Base.Type{Arrow.Interval{U, T}}) where {U, T} = Interval("interval", U == Arrow.Meta.IntervalUnit.YEAR_MONTH ? "YEAR_MONTH" : "DAY_TIME")
+StructTypes.StructType(::Base.Type{Interval}) = StructTypes.Struct()
+juliatype(f, x::Interval) = Arrow.Interval{x.unit == "YEAR_MONTH" ? Arrow.Meta.IntervalUnit.YEAR_MONTH : Arrow.Meta.IntervalUnit.DAY_TIME, x.unit == "YEAR_MONTH" ? Int32 : Int64}
+
+struct UnionT <: Type
+ name::String
+ mode::String
+ typIds::Vector{Int64}
+end
+
+Type(::Base.Type{Arrow.UnionT{T, typeIds, U}}) where {T, typeIds, U} = UnionT("union", T == Arrow.Meta.UnionMode.Dense ? "DENSE" : "SPARSE", collect(typeIds))
+children(::Base.Type{Arrow.UnionT{T, typeIds, U}}) where {T, typeIds, U} = Field[Field("", fieldtype(U, i), nothing) for i = 1:fieldcount(U)]
+StructTypes.StructType(::Base.Type{UnionT}) = StructTypes.Struct()
+juliatype(f, x::UnionT) = Arrow.UnionT{x.mode == "DENSE" ? Arrow.Meta.UnionMode.DENSE : Arrow.Meta.UnionMode.SPARSE, Tuple(x.typeIds), Tuple{(juliatype(y) for y in f.children)...}}
+
+struct List <: Type
+ name::String
+end
+
+Type(::Base.Type{Vector{T}}) where {T} = List("list")
+children(::Base.Type{Vector{T}}) where {T} = [Field("item", T, nothing)]
+StructTypes.StructType(::Base.Type{List}) = StructTypes.Struct()
+juliatype(f, x::List) = Vector{juliatype(f.children[1])}
+
+struct LargeList <: Type
+ name::String
+end
+
+StructTypes.StructType(::Base.Type{LargeList}) = StructTypes.Struct()
+juliatype(f, x::LargeList) = Vector{juliatype(f.children[1])}
+
+struct FixedSizeList <: Type
+ name::String
+ listSize::Int64
+end
+
+Type(::Base.Type{NTuple{N, T}}) where {N, T} = FixedSizeList("fixedsizelist", N)
+children(::Base.Type{NTuple{N, T}}) where {N, T} = [Field("item", T, nothing)]
+StructTypes.StructType(::Base.Type{FixedSizeList}) = StructTypes.Struct()
+juliatype(f, x::FixedSizeList) = NTuple{x.listSize, juliatype(f.children[1])}
+
+struct Struct <: Type
+ name::String
+end
+
+Type(::Base.Type{NamedTuple{names, types}}) where {names, types} = Struct("struct")
+children(::Base.Type{NamedTuple{names, types}}) where {names, types} = [Field(names[i], fieldtype(types, i), nothing) for i = 1:length(names)]
+StructTypes.StructType(::Base.Type{Struct}) = StructTypes.Struct()
+juliatype(f, x::Struct) = NamedTuple{Tuple(Symbol(x.name) for x in f.children), Tuple{(juliatype(y) for y in f.children)...}}
+
+struct Map <: Type
+ name::String
+ keysSorted::Base.Bool
+end
+
+Type(::Base.Type{Dict{K, V}}) where {K, V} = Map("map", false)
+children(::Base.Type{Dict{K, V}}) where {K, V} = [Field("entries", Arrow.KeyValue{K, V}, nothing)]
+StructTypes.StructType(::Base.Type{Map}) = StructTypes.Struct()
+juliatype(f, x::Map) = Dict{juliatype(f.children[1].children[1]), juliatype(f.children[1].children[2])}
+
+Type(::Base.Type{Arrow.KeyValue{K, V}}) where {K, V} = Struct("struct")
+children(::Base.Type{Arrow.KeyValue{K, V}}) where {K, V} = [Field("key", K, nothing), Field("value", V, nothing)]
+
+struct Null <: Type
+ name::String
+end
+
+Type(::Base.Type{Missing}) = Null("null")
+StructTypes.StructType(::Base.Type{Null}) = StructTypes.Struct()
+juliatype(f, x::Null) = Missing
+
+struct Utf8 <: Type
+ name::String
+end
+
+Type(::Base.Type{<:String}) = Utf8("utf8")
+StructTypes.StructType(::Base.Type{Utf8}) = StructTypes.Struct()
+juliatype(f, x::Utf8) = String
+
+struct LargeUtf8 <: Type
+ name::String
+end
+
+StructTypes.StructType(::Base.Type{LargeUtf8}) = StructTypes.Struct()
+juliatype(f, x::LargeUtf8) = String
+
+struct Binary <: Type
+ name::String
+end
+
+Type(::Base.Type{Vector{UInt8}}) = Binary("binary")
+children(::Base.Type{Vector{UInt8}}) = Field[]
+StructTypes.StructType(::Base.Type{Binary}) = StructTypes.Struct()
+juliatype(f, x::Binary) = Vector{UInt8}
+
+struct LargeBinary <: Type
+ name::String
+end
+
+StructTypes.StructType(::Base.Type{LargeBinary}) = StructTypes.Struct()
+juliatype(f, x::LargeBinary) = Vector{UInt8}
+
+struct Bool <: Type
+ name::String
+end
+
+Type(::Base.Type{Base.Bool}) = Bool("bool")
+StructTypes.StructType(::Base.Type{Bool}) = StructTypes.Struct()
+juliatype(f, x::Bool) = Base.Bool
+
+StructTypes.subtypekey(::Base.Type{Type}) = :name
+
+const SUBTYPES = @eval (
+ int=Int,
+ floatingpoint=FloatingPoint,
+ fixedsizebinary=FixedSizeBinary,
+ decimal=Decimal,
+ timestamp=Timestamp,
+ duration=Duration,
+ date=Date,
+ time=Time,
+ interval=Interval,
+ union=UnionT,
+ list=List,
+ largelist=LargeList,
+ fixedsizelist=FixedSizeList,
+ $(Symbol("struct"))=Struct,
+ map=Map,
+ null=Null,
+ utf8=Utf8,
+ largeutf8=LargeUtf8,
+ binary=Binary,
+ largebinary=LargeBinary,
+ bool=Bool
+)
+
+StructTypes.subtypes(::Base.Type{Type}) = SUBTYPES
+
+const Metadata = Union{Nothing, Vector{NamedTuple{(:key, :value), Tuple{String, String}}}}
+Metadata() = nothing
+
+mutable struct DictEncoding
+ id::Int64
+ indexType::Type
+ isOrdered::Base.Bool
+end
+
+DictEncoding() = DictEncoding(0, Type(), false)
+StructTypes.StructType(::Base.Type{DictEncoding}) = StructTypes.Mutable()
+
+mutable struct Field
+ name::String
+ nullable::Base.Bool
+ type::Type
+ children::Vector{Field}
+ dictionary::Union{DictEncoding, Nothing}
+ metadata::Metadata
+end
+
+Field() = Field("", true, Type(), Field[], nothing, Metadata())
+StructTypes.StructType(::Base.Type{Field}) = StructTypes.Mutable()
+Base.copy(f::Field) = Field(f.name, f.nullable, f.type, f.children, f.dictionary, f.metadata)
+
+function juliatype(f::Field)
+ T = juliatype(f, f.type)
+ return f.nullable ? Union{T, Missing} : T
+end
+
+function Field(nm, ::Base.Type{T}, dictencodings) where {T}
+ S = Arrow.maybemissing(T)
+ type = Type(S)
+ ch = children(S)
+ if dictencodings !== nothing && haskey(dictencodings, nm)
+ dict = dictencodings[nm]
+ else
+ dict = nothing
+ end
+ return Field(nm, T !== S, type, ch, dict, nothing)
+end
+
+mutable struct Schema
+ fields::Vector{Field}
+ metadata::Metadata
+end
+
+Schema() = Schema(Field[], Metadata())
+StructTypes.StructType(::Base.Type{Schema}) = StructTypes.Mutable()
+
+struct Offsets{T} <: AbstractVector{T}
+ data::Vector{T}
+end
+
+Base.size(x::Offsets) = size(x.data)
+Base.getindex(x::Offsets, i::Base.Int) = getindex(x.data, i)
+
+mutable struct FieldData
+ name::String
+ count::Int64
+ VALIDITY::Union{Nothing, Vector{Int8}}
+ OFFSET::Union{Nothing, Offsets}
+ TYPE_ID::Union{Nothing, Vector{Int8}}
+ DATA::Union{Nothing, Vector{Any}}
+ children::Vector{FieldData}
+end
+
+FieldData() = FieldData("", 0, nothing, nothing, nothing, nothing, FieldData[])
+StructTypes.StructType(::Base.Type{FieldData}) = StructTypes.Mutable()
+
+function FieldData(nm, ::Base.Type{T}, col, dictencodings) where {T}
+ if dictencodings !== nothing && haskey(dictencodings, nm)
+ refvals = DataAPI.refarray(col.data)
+ if refvals !== col.data
+ IT = eltype(refvals)
+ col = (x - one(T) for x in refvals)
+ else
+ _, de = dictencodings[nm]
+ IT = de.indexType
+ vals = unique(col)
+ col = Arrow.DictEncoder(col, vals, Arrow.encodingtype(length(vals)))
+ end
+ return FieldData(nm, IT, col, nothing)
+ end
+ S = Arrow.maybemissing(T)
+ len = Arrow._length(col)
+ VALIDITY = OFFSET = TYPE_ID = DATA = nothing
+ children = FieldData[]
+ if S <: Pair
+ return FieldData(nm, Vector{Arrow.KeyValue{Arrow._keytype(S), Arrow._valtype(S)}}, (Arrow.KeyValue(k, v) for (k, v) in pairs(col)))
+ elseif S !== Missing
+ # VALIDITY
+ VALIDITY = Int8[!ismissing(x) for x in col]
+ # OFFSET
+ if S <: Vector || S == String
+ lenfun = S == String ? x->ismissing(x) ? 0 : sizeof(x) : x->ismissing(x) ? 0 : length(x)
+ tot = sum(lenfun, col)
+ if tot > 2147483647
+ OFFSET = String[String(lenfun(x)) for x in col]
+ pushfirst!(OFFSET, "0")
+ else
+ OFFSET = Int32[ismissing(x) ? 0 : lenfun(x) for x in col]
+ pushfirst!(OFFSET, 0)
+ end
+ OFFSET = Offsets(OFFSET)
+ push!(children, FieldData("item", eltype(S), Arrow.flatten(skipmissing(col)), dictencodings))
+ elseif S <: NTuple
+ if Arrow.ArrowTypes.gettype(S) == UInt8
+ DATA = [ismissing(x) ? Arrow.ArrowTypes.default(S) : String(collect(x)) for x in col]
+ else
+ push!(children, FieldData("item", Arrow.ArrowTypes.gettype(S), Arrow.flatten(coalesce(x, Arrow.ArrowTypes.default(S)) for x in col), dictencodings))
+ end
+ elseif S <: NamedTuple
+ for (nm, typ) in zip(fieldnames(S), fieldtypes(S))
+ push!(children, FieldData(String(nm), typ, (getfield(x, nm) for x in col), dictencodings))
+ end
+ elseif S <: Arrow.UnionT
+ U = eltype(S)
+ tids = Arrow.typeids(S) === nothing ? (0:fieldcount(U)) : Arrow.typeids(S)
+ TYPE_ID = [x === missing ? 0 : tids[Arrow.isatypeid(x, U)] for x in col]
+ if Arrow.unionmode(S) == Arrow.Meta.UnionMode.Dense
+ offs = zeros(Int32, fieldcount(U))
+ OFFSET = Int32[]
+ for x in col
+ idx = x === missing ? 1 : Arrow.isatypeid(x, U)
+ push!(OFFSET, offs[idx])
+ offs[idx] += 1
+ end
+ for i = 1:fieldcount(U)
+ SS = fieldtype(U, i)
+ push!(children, FieldData("$i", SS, Arrow.filtered(i == 1 ? Union{SS, Missing} : Arrow.maybemissing(SS), col), dictencodings))
+ end
+ else
+ for i = 1:fieldcount(U)
+ SS = fieldtype(U, i)
+ push!(children, FieldData("$i", SS, Arrow.replaced(SS, col), dictencodings))
+ end
+ end
+ elseif S <: KeyValue
+ push!(children, FieldData("key", Arrow.keyvalueK(S), (x.key for x in col), dictencodings))
+ push!(children, FieldData("value", Arrow.keyvalueV(S), (x.value for x in col), dictencodings))
+ end
+ end
+ return FieldData(nm, len, VALIDITY, OFFSET, TYPE_ID, DATA, children)
+end
+
+mutable struct RecordBatch
+ count::Int64
+ columns::Vector{FieldData}
+end
+
+RecordBatch() = RecordBatch(0, FieldData[])
+StructTypes.StructType(::Base.Type{RecordBatch}) = StructTypes.Mutable()
+
+mutable struct DictionaryBatch
+ id::Int64
+ data::RecordBatch
+end
+
+DictionaryBatch() = DictionaryBatch(0, RecordBatch())
+StructTypes.StructType(::Base.Type{DictionaryBatch}) = StructTypes.Mutable()
+
+mutable struct DataFile <: Tables.AbstractColumns
+ schema::Schema
+ batches::Vector{RecordBatch}
+ dictionaries::Vector{DictionaryBatch}
+end
+
+Base.propertynames(x::DataFile) = (:schema, :batches, :dictionaries)
+
+function Base.getproperty(df::DataFile, nm::Symbol)
+ if nm === :schema
+ return getfield(df, :schema)
+ elseif nm === :batches
+ return getfield(df, :batches)
+ elseif nm === :dictionaries
+ return getfield(df, :dictionaries)
+ end
+ return Tables.getcolumn(df, nm)
+end
+
+DataFile() = DataFile(Schema(), RecordBatch[], DictionaryBatch[])
+StructTypes.StructType(::Base.Type{DataFile}) = StructTypes.Mutable()
+
+parsefile(file) = JSON3.read(Mmap.mmap(file), DataFile)
+
+# make DataFile satisfy Tables.jl interface
+function Tables.partitions(x::DataFile)
+ if isempty(x.batches)
+ # special case empty batches by producing a single DataFile w/ schema
+ return (DataFile(x.schema, RecordBatch[], x.dictionaries),)
+ else
+ return (DataFile(x.schema, [x.batches[i]], x.dictionaries) for i = 1:length(x.batches))
+ end
+end
+
+Tables.columns(x::DataFile) = x
+
+function Tables.schema(x::DataFile)
+ names = map(x -> x.name, x.schema.fields)
+ types = map(x -> juliatype(x), x.schema.fields)
+ return Tables.Schema(names, types)
+end
+
+Tables.columnnames(x::DataFile) = map(x -> Symbol(x.name), x.schema.fields)
+
+function Tables.getcolumn(x::DataFile, i::Base.Int)
+ field = x.schema.fields[i]
+ type = juliatype(field)
+ return ChainedVector(ArrowArray{type}[ArrowArray{type}(field, length(x.batches) > 0 ? x.batches[j].columns[i] : FieldData(), x.dictionaries) for j = 1:length(x.batches)])
+end
+
+function Tables.getcolumn(x::DataFile, nm::Symbol)
+ i = findfirst(x -> x.name == String(nm), x.schema.fields)
+ return Tables.getcolumn(x, i)
+end
+
+struct ArrowArray{T} <: AbstractVector{T}
+ field::Field
+ fielddata::FieldData
+ dictionaries::Vector{DictionaryBatch}
+end
+ArrowArray(f::Field, fd::FieldData, d) = ArrowArray{juliatype(f)}(f, fd, d)
+Base.size(x::ArrowArray) = (x.fielddata.count,)
+
+function Base.getindex(x::ArrowArray{T}, i::Base.Int) where {T}
+ @boundscheck checkbounds(x, i)
+ S = Base.nonmissingtype(T)
+ if x.field.dictionary !== nothing
+ fielddata = x.dictionaries[findfirst(y -> y.id == x.field.dictionary.id, x.dictionaries)].data.columns[1]
+ field = copy(x.field)
+ field.dictionary = nothing
+ idx = x.fielddata.DATA[i] + 1
+ return ArrowArray(field, fielddata, x.dictionaries)[idx]
+ end
+ if T === Missing
+ return missing
+ elseif S <: UnionT
+ U = eltype(S)
+ tids = Arrow.typeids(S) === nothing ? (0:fieldcount(U)) : Arrow.typeids(S)
+ typeid = tids[x.fielddata.TYPE_ID[i]]
+ if Arrow.unionmode(S) == Arrow.Meta.UnionMode.DENSE
+ off = x.fielddata.OFFSET[i]
+ return ArrowArray(x.field.children[typeid+1], x.fielddata.children[typeid+1], x.dictionaries)[off]
+ else
+ return ArrowArray(x.field.children[typeid+1], x.fielddata.children[typeid+1], x.dictionaries)[i]
+ end
+ end
+ x.fielddata.VALIDITY[i] == 0 && return missing
+ if S <: Vector{UInt8}
+ return copy(codeunits(x.fielddata.DATA[i]))
+ elseif S <: String
+ return x.fielddata.DATA[i]
+ elseif S <: Vector
+ offs = x.fielddata.OFFSET
+ A = ArrowArray{eltype(S)}(x.field.children[1], x.fielddata.children[1], x.dictionaries)
+ return A[(offs[i] + 1):offs[i + 1]]
+ elseif S <: Dict
+ offs = x.fielddata.OFFSET
+ A = ArrowArray(x.field.children[1], x.fielddata.children[1], x.dictionaries)
+ return Dict(y.key => y.value for y in A[(offs[i] + 1):offs[i + 1]])
+ elseif S <: Tuple
+ if Arrow.ArrowTypes.gettype(S) == UInt8
+ A = x.fielddata.DATA
+ return Tuple(map(UInt8, collect(A[i][1:x.field.type.byteWidth])))
+ else
+ sz = x.field.type.listSize
+ A = ArrowArray{Arrow.ArrowTypes.gettype(S)}(x.field.children[1], x.fielddata.children[1], x.dictionaries)
+ off = (i - 1) * sz + 1
+ return Tuple(A[off:(off + sz - 1)])
+ end
+ elseif S <: NamedTuple
+ data = (ArrowArray(x.field.children[j], x.fielddata.children[j], x.dictionaries)[i] for j = 1:length(x.field.children))
+ return NamedTuple{fieldnames(S)}(Tuple(data))
+ elseif S == Int64 || S == UInt64
+ return parse(S, x.fielddata.DATA[i])
+ elseif S <: Arrow.Decimal
+ str = x.fielddata.DATA[i]
+ return S(parse(Int128, str))
+ elseif S <: Arrow.Date || S <: Arrow.Time
+ val = x.fielddata.DATA[i]
+ return Arrow.storagetype(S) == Int32 ? S(val) : S(parse(Int64, val))
+ elseif S <: Arrow.Timestamp
+ return S(parse(Int64, x.fielddata.DATA[i]))
+ else
+ return S(x.fielddata.DATA[i])
+ end
+end
+
+# take any Tables.jl source and write out arrow json datafile
+function DataFile(source)
+ fields = Field[]
+ metadata = nothing # TODO?
+ batches = RecordBatch[]
+ dictionaries = DictionaryBatch[]
+ dictencodings = Dict{String, Tuple{Base.Type, DictEncoding}}()
+ dictid = Ref(0)
+ for (i, tbl1) in Tables.partitions(source)
+ tbl = Arrow.toarrowtable(tbl1)
+ if i == 1
+ sch = Tables.schema(tbl)
+ for (nm, T, col) in zip(sch.names, sch.types, Tables.Columns(tbl))
+ if col isa Arrow.DictEncode
+ id = dictid[]
+ dictid[] += 1
+ codes = DataAPI.refarray(col.data)
+ if codes !== col.data
+ IT = Type(eltype(codes))
+ else
+ IT = Type(Arrow.encodingtype(length(unique(col))))
+ end
+ dictencodings[String(nm)] = (T, DictEncoding(id, IT, false))
+ end
+ push!(fields, Field(String(nm), T, dictencodings))
+ end
+ end
+ # build record batch
+ len = Tables.rowcount(tbl)
+ columns = FieldData[]
+ for (nm, T, col) in zip(sch.names, sch.types, Tables.Columns(tbl))
+ push!(columns, FieldData(String(nm), T, col, dictencodings))
+ end
+ push!(batches, RecordBatch(len, columns))
+ # build dictionaries
+ for (nm, (T, dictencoding)) in dictencodings
+ column = FieldData(nm, T, Tables.getcolumn(tbl, nm), nothing)
+ recordbatch = RecordBatch(len, [column])
+ push!(dictionaries, DictionaryBatch(dictencoding.id, recordbatch))
+ end
+ end
+ schema = Schema(fields, metadata)
+ return DataFile(schema, batches, dictionaries)
+end
+
+function Base.isequal(df::DataFile, tbl::Arrow.Table)
+ Tables.schema(df) == Tables.schema(tbl) || return false
+ i = 1
+ for (col1, col2) in zip(Tables.Columns(df), Tables.Columns(tbl))
+ if !isequal(col1, col2)
+ @show i
+ return false
+ end
+ i += 1
+ end
+ return true
+end
+
+end
projects / ceph.git / blobdiff