Skip to content

Commit 08bc67b

Browse files
iblislinjpsamaroo
authored andcommitted
Add blocked CSV reading functionality
Import BlockIO/ChunkIter from Dagger Wire blocking into loadtable
1 parent fff0206 commit 08bc67b

4 files changed

Lines changed: 181 additions & 6 deletions

File tree

Project.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
name = "JuliaDB"
22
uuid = "a93385a2-3734-596a-9a66-3cfbb77141e6"
3+
version = "0.12.0"
34

45
[deps]
56
Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
67
DataValues = "e7dc6d0d-1eca-5fa6-8ad6-5aecde8b7ea5"
8+
DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
79
Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
810
Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
911
IndexedTables = "6deec6e2-d858-57c5-ab9b-e6ca5bd20e43"

src/JuliaDB.jl

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ Base.isequal(x::Dataset, y::DDataset) = isequal(x, collect(y))
6666
include("iteration.jl")
6767
include("sort.jl")
6868

69+
include("block-io.jl")
6970
include("io.jl")
7071
include("printing.jl")
7172

src/block-io.jl

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
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], ',')

src/util.jl

Lines changed: 32 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,9 +54,9 @@ function _loadtable_serial(T, file::Union{IO, AbstractString, AbstractArray};
5454
presorted=false,
5555
copy=false,
5656
csvread=TextParse.csvread,
57+
nblocks=nworkers(),
5758
kwargs...)
5859

59-
#println("LOADING ", file)
6060
count = Int[]
6161

6262
samecols = nothing
@@ -79,12 +79,38 @@ function _loadtable_serial(T, file::Union{IO, AbstractString, AbstractArray};
7979
samecols = map(x->map(string, x), samecols)
8080
end
8181

82-
if isa(file, AbstractArray)
83-
cols, header, count = csvread(file, delim;
84-
samecols=samecols,
85-
kwargs...)
82+
if nblocks > 1
83+
if !(file isa AbstractArray)
84+
file = [file,]
85+
end
86+
count = length(file)
87+
cols = nothing
88+
header = nothing
89+
for f in file
90+
bios = blocks(f, '\n', nblocks)
91+
for (idx,b) in enumerate(bios)
92+
_cols, _header = csvread(b, delim;
93+
header_exists=(idx==1),
94+
kwargs...)
95+
if header === nothing
96+
header = _header
97+
cols = _cols
98+
else
99+
@assert length(_cols) == length(cols)
100+
for idx in 1:length(cols)
101+
append!(cols[idx], _cols[idx])
102+
end
103+
end
104+
end
105+
end
86106
else
87-
cols, header = csvread(file, delim; kwargs...)
107+
if file isa AbstractArray
108+
cols, header, count = csvread(file, delim;
109+
samecols=samecols,
110+
kwargs...)
111+
else
112+
cols, header = csvread(file, delim; kwargs...)
113+
end
88114
end
89115

90116
header = map(string, header)

0 commit comments

Comments
 (0)