Skip to content

Commit affcd1b

Browse files
guangtaoguangtao
guangtao
authored and
guangtao
committed
Tighten Flight test hygiene and ArrowTypes support
1 parent d9b4bc9 commit affcd1b

File tree

91 files changed

+1581
-37
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

91 files changed

+1581
-37
lines changed

‎docs/src/manual.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,10 +97,20 @@ One note on performance: when writing `TimeZones.ZonedDateTime` columns to the a
9797
as the column has `ZonedDateTime` elements that all share a common timezone. This ensures the writing process can know "upfront" which timezone will be encoded and is thus much more
9898
efficient and performant.
9999

100+
Similarly, `ArrowTypes.ToArrow` avoids repeated type-promotion work for
101+
homogeneous custom columns even when `ArrowTypes.ArrowType(T)` is abstract, so
102+
write-time conversion does not pay unnecessary overhead once the serialized
103+
element type is stable.
104+
100105
#### Custom types
101106

102107
To support writing your custom Julia struct, Arrow.jl utilizes the format's mechanism for "extension types" by allowing the storing of Julia type name and metadata in the field metadata. To "hook in" to this machinery, custom types can utilize the interface methods defined in the `Arrow.ArrowTypes` submodule. For example:
103108

109+
Arrow.jl already uses this mechanism for several Base logical types, including
110+
`nothing`, `Tuple`, `VersionNumber`, and `Complex`, so those values roundtrip as
111+
their original Julia types instead of falling back to plain struct-shaped
112+
`NamedTuple`s.
113+
104114
```julia
105115
using Arrow
106116

‎ext/ArrowgRPCServerExt.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
module ArrowgRPCServerExt
219

320
using Arrow

‎ext/arrowgrpcserverext/constants.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
const Flight = Arrow.Flight
219
const STREAM_BUFFER_SIZE = 16
320
const GENERATED_TYPE_PREFIX = "Arrow.Flight.Generated."

‎ext/arrowgrpcserverext/context.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
function _method_type(method::Flight.MethodDescriptor)
219
if method.request_streaming
320
return method.response_streaming ? gRPCServer.MethodType.BIDI_STREAMING :

‎ext/arrowgrpcserverext/descriptor.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
function _register_proto_types!(method::Flight.MethodDescriptor)
219
registry = gRPCServer.get_type_registry()
320
registry[_proto_type_name(method.request_type)] = method.request_type

‎ext/arrowgrpcserverext/handlers.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
function _unary_handler(service::Flight.Service, method::Flight.MethodDescriptor)
219
return (context, request) -> Flight.dispatch(service, _call_context(context), method, request)
320
end

‎ext/arrowgrpcserverext/streams.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
function _drain_response!(
219
stream::gRPCServer.ServerStream{T},
320
response::Channel{T},

‎src/ArrowTypes/src/ArrowTypes.jl‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,14 @@ arrowname(::Type{Tuple{}}) = TUPLE
324324
JuliaType(::Val{TUPLE}, ::Type{NamedTuple{names,types}}) where {names,types<:Tuple} = types
325325
fromarrow(::Type{T}, x::NamedTuple) where {T<:Tuple} = Tuple(x)
326326

327+
# Complex
328+
const COMPLEX = Symbol("JuliaLang.Complex")
329+
arrowname(::Type{<:Complex}) = COMPLEX
330+
JuliaType(::Val{COMPLEX}, ::Type{NamedTuple{names,Tuple{T,T}}}) where {names,T<:Real} =
331+
Complex{T}
332+
fromarrowstruct(::Type{T}, ::Val{(:re, :im)}, re, im) where {T<:Complex} = T(re, im)
333+
fromarrowstruct(::Type{T}, ::Val{(:im, :re)}, im, re) where {T<:Complex} = T(re, im)
334+
327335
# VersionNumber
328336
const VERSION_NUMBER = Symbol("JuliaLang.VersionNumber")
329337
ArrowKind(::Type{VersionNumber}) = StructKind()
@@ -388,10 +396,26 @@ default(::Type{NamedTuple{names,types}}) where {names,types} =
388396
NamedTuple{names}(Tuple(default(fieldtype(types, i)) for i = 1:length(names)))
389397

390398
function promoteunion(T, S)
399+
T === S && return T
391400
new = promote_type(T, S)
392401
return isabstracttype(new) ? Union{T,S} : new
393402
end
394403

404+
function _toarroweltype(x)
405+
state = iterate(x)
406+
state === nothing && return Missing
407+
y, st = state
408+
T = typeof(toarrow(y))
409+
while true
410+
state = iterate(x, st)
411+
state === nothing && return T
412+
y, st = state
413+
S = typeof(toarrow(y))
414+
S === T && continue
415+
T = promoteunion(T, S)
416+
end
417+
end
418+
395419
# lazily call toarrow(x) on getindex for each x in data
396420
struct ToArrow{T,A} <: AbstractVector{T}
397421
data::A
@@ -412,7 +436,7 @@ function ToArrow(x::A) where {A}
412436
if isempty(x)
413437
return Missing[]
414438
end
415-
T = mapreduce(typeof ∘ toarrow, promoteunion, x)
439+
T = _toarroweltype(x)
416440
if T === Missing && concrete_or_concreteunion(S)
417441
T = promoteunion(T, typeof(toarrow(default(S))))
418442
end

‎src/ArrowTypes/test/tests.jl‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -144,6 +144,17 @@ end
144144
@test ArrowTypes.default(Tuple{Vararg{Int}}) == ()
145145
@test ArrowTypes.default(Tuple{String,Vararg{Int}}) == ("",)
146146

147+
z = 1.0 + 2.0im
148+
@test ArrowTypes.ArrowKind(typeof(z)) == ArrowTypes.StructKind()
149+
@test ArrowTypes.arrowname(typeof(z)) == ArrowTypes.COMPLEX
150+
@test ArrowTypes.arrowname(Union{Missing,typeof(z)}) == ArrowTypes.COMPLEX
151+
@test ArrowTypes.JuliaType(
152+
Val(ArrowTypes.COMPLEX),
153+
NamedTuple{(:re, :im),Tuple{Float64,Float64}},
154+
) == ComplexF64
155+
@test ArrowTypes.fromarrowstruct(ComplexF64, Val((:re, :im)), 1.0, 2.0) == z
156+
@test ArrowTypes.fromarrowstruct(ComplexF64, Val((:im, :re)), 2.0, 1.0) == z
157+
147158
v = v"1"
148159
v_nt = (major=1, minor=0, patch=0, prerelease=(), build=())
149160
@test ArrowTypes.ArrowKind(VersionNumber) == ArrowTypes.StructKind()
@@ -167,6 +178,7 @@ end
167178

168179
@test ArrowTypes.promoteunion(Int, Float64) == Float64
169180
@test ArrowTypes.promoteunion(Int, String) == Union{Int,String}
181+
@test ArrowTypes.promoteunion(Int, Int) == Int
170182

171183
@test ArrowTypes.concrete_or_concreteunion(Int)
172184
@test !ArrowTypes.concrete_or_concreteunion(Union{Real,String})
@@ -219,6 +231,12 @@ end
219231
T = Union{DateTimeTZ,Missing}
220232
@test !ArrowTypes.concrete_or_concreteunion(ArrowTypes.ArrowType(T))
221233
@test eltype(ArrowTypes.ToArrow(T[missing])) == Union{Timestamp{:UTC},Missing}
234+
@test eltype(ArrowTypes.ToArrow(DateTimeTZ[DateTimeTZ(1, "UTC"), DateTimeTZ(2, "UTC")])) ==
235+
Timestamp{:UTC}
236+
@test eltype(ArrowTypes.ToArrow(DateTimeTZ[DateTimeTZ(1, "UTC"), DateTimeTZ(2, "PST")])) ==
237+
Timestamp
238+
@test eltype(ArrowTypes.ToArrow(Any[DateTimeTZ(1, "UTC"), DateTimeTZ(2, "UTC")])) ==
239+
Timestamp{:UTC}
222240

223241
# Works since `ArrowTypes.default(Any) === nothing` and
224242
# `ArrowTypes.toarrow(nothing) === missing`. Defining `toarrow(::Nothing) = nothing`

‎src/flight/Flight.jl‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,20 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
118
module Flight
219

320
using Base64

0 commit comments

Comments
 (0)