Skip to content

Commit

Permalink
Admit arbitrary ProtoType for read and write
Browse files Browse the repository at this point in the history
  • Loading branch information
lassepe committed Mar 18, 2021
1 parent 95cc27b commit 8774420
Showing 1 changed file with 10 additions and 8 deletions.
18 changes: 10 additions & 8 deletions src/core.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ using Base.Threads
using CodecZlib
using BufferedStreams
using MacroTools: @forward
using ProtoBuf: ProtoType

# Ref: https://github.com/tensorflow/tensorflow/blob/295ad2781683835be974faba0a191528d8079768/tensorflow/core/lib/hash/crc32c.h#L50-L59

Expand Down Expand Up @@ -62,10 +63,10 @@ read(f, s::String; kwargs...) = read(f, [s]; kw...)
function read(
f,
files::Vector;
compression=nothing,
bufsize=10*1024*1024,
channel_size=1_000,
record_type=Example
compression = nothing,
bufsize = 10 * 1024 * 1024,
channel_size = 1_000,
record_type = Example,
)
Channel{record_type}(channel_size) do ch
@threads for file_name in files
Expand All @@ -77,12 +78,13 @@ function read(
elseif compression == :zlib
io = ZlibDecompressorStream(io)
else
isnothing(compression) || throw(ArgumentError("unsupported compression method: $compression"))
isnothing(compression) ||
throw(ArgumentError("unsupported compression method: $compression"))
end

while !eof(io)
example = readproto(IOBuffer(read_record(io)), Example())
put!(ch, f(example))
instance = readproto(IOBuffer(read_record(io)), record_type())
put!(ch, f(instance))
end
end
end
Expand Down Expand Up @@ -128,7 +130,7 @@ function write(io::IO, xs)
end
end

function write(io::IO, x::Example)
function write(io::IO, x::ProtoType)
buff = IOBuffer()
writeproto(buff, x)

Expand Down

0 comments on commit 8774420

Please sign in to comment.