Source file owl_io.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
# 1 "src/base/misc/owl_io.ml"
(*
 * OWL - OCaml Scientific and Engineering Computing
 * Copyright (c) 2016-2020 Liang Wang <liang.wang@cl.cam.ac.uk>
 *)

open Owl_utils

(* read a file of a given path *)
let read_file ?(trim = true) f =
  let h = open_in f in
  let s = Stack.make () in
  (try
     while true do
       let l =
         match trim with
         | true  -> input_line h |> String.trim
         | false -> input_line h
       in
       Stack.push s l
     done
   with
  | End_of_file -> ());
  close_in h;
  Stack.to_array s


let read_file_string f =
  let ic = open_in f in
  let n = in_channel_length ic in
  let s = Bytes.create n in
  really_input ic s 0 n;
  close_in ic;
  Bytes.to_string s


(* write a file of a given path *)
let write_file ?(_flag = Open_creat) f s =
  let h = open_out f in
  Printf.fprintf h "%s" s;
  close_out h


(* iterate every doc in the corpus without loading the whole corpus in the
  memory, then apply passed in function f. note that each line is a doc.
 *)
let iteri_lines_of_file ?(verbose = true) f fname =
  let i = ref 0 in
  let h = open_in fname in
  (let t0 = Unix.gettimeofday () in
   let t1 = ref (Unix.gettimeofday ()) in
   try
     while true do
       f !i (input_line h);
       i := !i + 1;
       (* output summary if in verbose mode *)
       if verbose = true
       then (
         let t2 = Unix.gettimeofday () in
         if t2 -. !t1 > 5.
         then (
           t1 := t2;
           let speed = float_of_int !i /. (t2 -. t0) |> int_of_float in
           Owl_log.info "processed %i, avg. %i docs/s" !i speed))
     done
   with
   | End_of_file -> ());
  close_in h


(* map every doc in the corpus into another type *)
let mapi_lines_of_file f fname =
  let stack = Owl_utils.Stack.make () in
  iteri_lines_of_file (fun i s -> Owl_utils.Stack.push stack (f i s)) fname;
  Owl_utils.Stack.to_array stack


(* similar to iteri_lines_of_file but for marshaled file *)
let iteri_lines_of_marshal ?(verbose = true) f fname =
  let i = ref 0 in
  let h = open_in fname in
  (let t1 = ref (Unix.gettimeofday ()) in
   let i1 = ref 0 in
   try
     while true do
       f !i (Marshal.from_channel h);
       i := !i + 1;
       (* output summary if in verbose mode *)
       if verbose = true
       then (
         let t2 = Unix.gettimeofday () in
         if t2 -. !t1 > 5.
         then (
           let speed = float_of_int (!i - !i1) /. (t2 -. !t1) |> int_of_float in
           i1 := !i;
           t1 := t2;
           Owl_log.info "processed %i, avg. %i docs/s" !i speed))
     done
   with
   | End_of_file -> ());
  close_in h


(* similar to mapi_lines_of_file but for marshaled file *)
let mapi_lines_of_marshal f fname =
  let stack = Owl_utils.Stack.make () in
  iteri_lines_of_marshal (fun i s -> Owl_utils.Stack.push stack (f i s)) fname;
  Owl_utils.Stack.to_array stack


(* save a marshalled object to a file *)
let marshal_to_file ?(flags = []) x f =
  let h = open_out f in
  Marshal.to_channel h x flags;
  close_out h


(* load a marshalled object from a file *)
let marshal_from_file f =
  let h = open_in f in
  let s = really_input_string h (in_channel_length h) in
  Marshal.from_string s 0


let head n fname =
  let lines = Owl_utils.Stack.make () in
  (try
     iteri_lines_of_file
       (fun i s ->
         assert (i < n);
         Owl_utils.Stack.push lines s)
       fname
   with
  | _exn -> ());
  Owl_utils.Stack.to_array lines


(* TODO *)
let _tail _n _fname = raise (Owl_exception.NOT_IMPLEMENTED "owl_io._tail")

let read_csv ?(sep = '\t') fname =
  let lines = Owl_utils.Stack.make () in
  iteri_lines_of_file
    (fun _i s ->
      String.trim s
      |> String.split_on_char sep
      |> Array.of_list
      |> Owl_utils.Stack.push lines)
    fname;
  Owl_utils.Stack.to_array lines


let write_csv ?(sep = '\t') x fname =
  let h = open_out fname in
  Array.iter
    (fun row ->
      let s =
        Array.fold_left (fun acc elt -> Printf.sprintf "%s%s%c" acc elt sep) "" row
      in
      Printf.fprintf h "%s\n" s)
    x;
  close_out h


let read_csv_proc ?(sep = '\t') proc fname =
  iteri_lines_of_file
    (fun i s -> String.trim s |> String.split_on_char sep |> Array.of_list |> proc i)
    fname


let write_csv_proc ?(sep = '\t') x proc fname =
  let h = open_out fname in
  Array.iter
    (fun row ->
      let s =
        Array.fold_left (fun acc elt -> Printf.sprintf "%s%s%c" acc (proc elt) sep) "" row
      in
      Printf.fprintf h "%s\n" s)
    x;
  close_out h


let csv_head ?(sep = '\t') idx fname =
  let h = open_in fname in
  for _i = 1 to idx - 1 do
    input_line h |> ignore
  done;
  input_line h |> String.trim |> String.split_on_char sep |> Array.of_list