|
| 1 | +## |
| 2 | +# give a stream view to a block from any seekable iostream |
| 3 | +# By @tanmaykm |
| 4 | +# |
| 5 | +import Base: close, eof, read, read!, peek, seek, write, filesize, position, seekend, seekstart, skip, bytesavailable |
| 6 | +using DelimitedFiles |
| 7 | + |
| 8 | +############################################################################### |
| 9 | +# BlockIO |
| 10 | +############################################################################### |
| 11 | + |
| 12 | +struct BlockIO <: IO |
| 13 | + s::IO |
| 14 | + r::UnitRange |
| 15 | + l::Int |
| 16 | + |
| 17 | + function find_end_pos(bio::BlockIO, end_byte::Char) |
| 18 | + seekend(bio) |
| 19 | + try |
| 20 | + while(!eof(bio.s) && (end_byte != read(bio, Char))) continue end |
| 21 | + catch |
| 22 | + end |
| 23 | + position(bio.s) |
| 24 | + end |
| 25 | + |
| 26 | + function find_start_pos(bio::BlockIO, end_byte::Char) |
| 27 | + (bio.r.start == 1) && (return bio.r.start) |
| 28 | + seekstart(bio) |
| 29 | + !eof(bio.s) && while(end_byte != read(bio, Char)) continue end |
| 30 | + position(bio.s)+1 |
| 31 | + end |
| 32 | + |
| 33 | + function BlockIO(s::IO, r::UnitRange, match_ends::Union{Char,Nothing}=nothing) |
| 34 | + # TODO: use mark when available |
| 35 | + seekend(s) |
| 36 | + ep = position(s) |
| 37 | + |
| 38 | + r = min(r.start,ep+1):min(r.start+length(r)-1,ep) |
| 39 | + bio = new(s, r, length(r)) |
| 40 | + if(nothing != match_ends) |
| 41 | + p1 = find_start_pos(bio, match_ends) |
| 42 | + p2 = find_end_pos(bio, match_ends) |
| 43 | + r = p1:p2 |
| 44 | + bio = new(s, r, length(r)) |
| 45 | + end |
| 46 | + seekstart(bio) |
| 47 | + bio |
| 48 | + end |
| 49 | +end |
| 50 | + |
| 51 | +BlockIO(bio::BlockIO, match_ends::Union{Char,Nothing}=nothing) = BlockIO(bio.s, bio.r, match_ends) |
| 52 | + |
| 53 | +close(bio::BlockIO) = close(bio.s) |
| 54 | +eof(bio::BlockIO) = (position(bio) >= bio.l) |
| 55 | +read(bio::BlockIO, x::Type{UInt8}) = read(bio.s, x) |
| 56 | +read!(bio::BlockIO, a::Vector{UInt8}) = (length(a) <= bytesavailable(bio)) ? read!(bio.s, a) : throw(EOFError()) |
| 57 | +read!(bio::BlockIO, a::Array{T}) where {T} = (length(a)*sizeof(T) <= bytesavailable(bio)) ? read!(bio.s, a) : throw(EOFError()) |
| 58 | + |
| 59 | +read(bio::BlockIO, nb::Integer = bio.l) = String(read!(bio, Array{UInt8}(undef, nb))) |
| 60 | + |
| 61 | +peek(bio::BlockIO) = peek(bio.s) |
| 62 | +write(bio::BlockIO, p::Ptr, nb::Integer) = write(bio, p, int(nb)) |
| 63 | +write(bio::BlockIO, p::Ptr, nb::Int) = write(bio.s, p, nb) |
| 64 | +write(bio::BlockIO, x::UInt8) = write(bio, UInt8[x]) |
| 65 | +write(bio::BlockIO, a::Array{T}, len) where {T} = write_sub(bio, a, 1, len) |
| 66 | +write(bio::BlockIO, a::Array{T}) where {T} = write(bio, a, length(a)) |
| 67 | +write_sub(bio::BlockIO, a::Array{T}, offs, len) where {T} = isbits(T) ? write(bio, pointer(a,offs), len*sizeof(T)) : error("$T is not bits type") |
| 68 | + |
| 69 | +bytesavailable(bio::BlockIO) = (bio.l - position(bio)) |
| 70 | +position(bio::BlockIO) = position(bio.s) - bio.r.start + 1 |
| 71 | + |
| 72 | +filesize(bio::BlockIO) = bio.l |
| 73 | + |
| 74 | +seek(bio::BlockIO, n::Integer) = seek(bio.s, n+bio.r.start-1) |
| 75 | +seekend(bio::BlockIO) = seek(bio, filesize(bio)) |
| 76 | +seekstart(bio::BlockIO) = seek(bio, 0) |
| 77 | +skip(bio::BlockIO, n::Integer) = seek(bio, n+position(bio)) |
| 78 | + |
| 79 | +# it returns a IOBuffer, not String |
| 80 | +function readuntil(bio::BlockIO, dlm::Char, n::Int) |
| 81 | + io = IOBuffer() |
| 82 | + i = 0 |
| 83 | + while !eof(bio) |
| 84 | + c = read(bio, Char) |
| 85 | + write(io, c) |
| 86 | + c == dlm && (i += 1) |
| 87 | + i == n && break |
| 88 | + end |
| 89 | + seekstart(io) |
| 90 | +end |
| 91 | + |
| 92 | +blocksize(s, n) = ceil(Int, s / n) |
| 93 | + |
| 94 | +function blocks(fname, dlm = '\n', n = length(workers())) |
| 95 | + s = filesize(fname) |
| 96 | + bs = blocksize(s, n) |
| 97 | + @assert(bs > 1, "Specified more blocks than bytes in file") |
| 98 | + |
| 99 | + start = 1 |
| 100 | + stop = 1 |
| 101 | + offset = 0 # bs offset |
| 102 | + bios = Vector{BlockIO}(undef, n) |
| 103 | + |
| 104 | + for i ∈ 1:n |
| 105 | + stop = start + offset + bs - 1 |
| 106 | + r = start:stop |
| 107 | + bio = BlockIO(open(fname), r, dlm) |
| 108 | + stop′ = bio.r.stop # get relocated ending position |
| 109 | + start = stop′ |
| 110 | + offset = stop′ - stop |
| 111 | + bios[i] = bio |
| 112 | + end |
| 113 | + |
| 114 | + @assert sum(filesize.(bios)) == s |
| 115 | + |
| 116 | + bios |
| 117 | +end |
| 118 | + |
| 119 | +############################################################################### |
| 120 | +# ChunkIter |
| 121 | +############################################################################### |
| 122 | + |
| 123 | +struct ChunkIter |
| 124 | + io::IO |
| 125 | + dlm::Char |
| 126 | + n::Int |
| 127 | +end |
| 128 | + |
| 129 | +ChunkIter(fname::String, dlm::Char = '\n', n::Int = 1000) = |
| 130 | + ChunkIter(open(fname), dlm, n) |
| 131 | + |
| 132 | +function Base.iterate(i::ChunkIter, idx=1) |
| 133 | + eof(i.io) && return nothing |
| 134 | + return readuntil(i.io, i.dlm, i.n), idx+1 |
| 135 | +end |
| 136 | + |
| 137 | +############################################################################### |
| 138 | +# Usage |
| 139 | +############################################################################### |
| 140 | + |
| 141 | +# addprocs() |
| 142 | +# fname = "path/to/file.csv" |
| 143 | +# bios = JuliaDB.blocks(fname) |
| 144 | +# b = bios[1] |
| 145 | +# iter = ChunkIter(b, '\n', 10) |
| 146 | +# readdlm(next(iter, 42)[1], ',') |
0 commit comments