-module(csv_reader). -on_load(init_nif/0). -define(D(X), io:format("~p:~p ~p~n", [?MODULE, ?LINE, X])). % -define(D(X), ok). -export([init/2, next/1, total_lines/1]). -export([date_to_ms/2]). -export([foldl/4]). foldl(Fun, InitialState, CSVPath, Options) when is_function(Fun) andalso is_list(CSVPath) -> {ok, Reader} = init(CSVPath, Options), Total = total_lines(Reader), State1 = Fun({init_total, Total}, InitialState), foldl_0(Fun, State1, Reader). foldl_0(Fun, State, Reader) -> case next(Reader) of {ok, Lines, Reader1} -> State1 = lists:foldl(Fun, State, [Line || Line <- Lines]), foldl_0(Fun, State1, Reader1); {eof, _} -> Fun(eof, State) end. date_to_ms({YY,MM,DD},{H,M,S,MS}) -> date_to_ms_nif(YY, MM, DD, H, M, S, MS); date_to_ms(_,_) -> undefined. date_to_ms_nif(YY, MM, DD, H, M, S, MS) -> Date = {YY, MM, DD}, Time = {H,M,S}, Timestamp = (calendar:datetime_to_gregorian_seconds({Date,Time}) - calendar:datetime_to_gregorian_seconds({{1970,1,1}, {0,0,0}})), Timestamp*1000 + MS. init_nif() -> Path = filename:dirname(code:which(?MODULE)) ++ "/../priv", Load = erlang:load_nif(Path ++ "/csv_reader", 0), case Load of ok -> ok; {error, {Reason,Text}} -> io:format("Load csv_reader failed. ~p:~p~n", [Reason, Text]) end, ok. -record(loader, { file, fd, options = [], offset, limit, header, pattern, cols, count = 0, total, buffer = <<>> }). parse_line(Bin, _Pattern) -> case binary:split(Bin, <<"\n">>) of [A,B] -> {A, B}; [A] -> {undefined, A} end. filter(float) -> $f; filter(int) -> $i; filter(date) -> $d; filter(time) -> $t; filter(utc) -> $g; filter(_) -> $u. total_lines(#loader{total = Total}) -> Total. compile_pattern(Cols1, Options) -> Cols = [binary_to_list(H) || H <- Cols1], Record = atom_to_binary(proplists:get_value(header, Options, csv), latin1), RecordSize = proplists:get_value(size, Options, length(Cols) + 1), Map = lists:map(fun(H) -> case lists:keyfind(H, 1, Options) of {H, OutPos, Filter} -> <>; false -> <<-1, $s>> end end, Cols), iolist_to_binary([<<(size(Record)), Record/binary, RecordSize, (length(Map))>>, Map]). init(Path, Options) -> SkipLines = lists:member(skip_lines, Options), {OpenOptions, TotalCmd} = case re:run(Path, "\\.gz$") of nomatch when SkipLines -> {[{read_ahead, 1024*1024}], "echo 0"}; _ when SkipLines -> {[compressed], "echo 0"}; nomatch -> {[{read_ahead, 1024*1024}], "wc -l \""++Path++"\""}; _ -> {[compressed], "gzcat \""++Path++"\" | wc -l"} end, {match, [Tot]} = re:run(os:cmd(TotalCmd), "(\\d+)", [{capture,all_but_first,list}]), Total = list_to_integer(Tot), ?D({skip_lines, SkipLines, Total}), case file:open(Path, [raw, binary|OpenOptions]) of {ok, F} -> start_loader1(Path, Options, F, Total); {error, Error} -> {error, Error} end. start_loader1(Path, Options, F, Total) -> {ok, Header1} = file:read_line(F), [Header2, <<>>] = binary:split(Header1, [<<"\n">>]), Header = binary:split(Header2, [<<",">>], [global]), Pattern = compile_pattern(Header, Options), Loader1 = #loader{ file = Path, header = Header, offset = size(Header1), cols = length(Header), pattern = Pattern, fd = F, total = Total, options = Options }, {ok, Loader1}. next(#loader{fd = F, buffer = Buffer, pattern = Pattern, count = Count} = Loader) -> case file:read(F, 65536) of {ok, Bin} -> {Lines, Rest} = split_lines(<>, Pattern), {ok, Lines, Loader#loader{buffer = Rest, count = Count + length(Lines)}}; {error, Error} -> {error, Error}; eof when Buffer == <<>> -> {eof, Count}; eof -> {Lines, _} = split_lines(<>, Pattern), {ok, Lines, Loader#loader{buffer = <<>>, count = Count + length(Lines)}} end. split_lines(Bin, Pattern) -> split_lines(Bin, [], Pattern). split_lines(Bin, Acc, Pattern) -> case parse_line(Bin, Pattern) of {undefined, Rest} -> {lists:reverse(Acc), Rest}; {Line, Rest} -> % ?D(Line), split_lines(Rest, [Line|Acc], Pattern) end.