tk/Dataframe: fixes to the API

lib/toolbox/dataframe.ml

@@ -2,12 +2,12 @@ open Core_kernel
 open Rresult
 
 type column =
-  | Int of int array
-  | Maybe_int of int option array
-  | Float of float array
-  | Maybe_float of float option array
-  | String of string array
-  | Maybe_string of string option array
+  | Ints of int array
+  | Int_opts of int option array
+  | Floats of float array
+  | Float_opts of float option array
+  | Strings of string array
+  | String_opts of string option array
 [@@deriving show]
 
 type t = {
@@ -19,7 +19,10 @@ type t = {
 let nrows t = t.nrows
 let ncols t = t.ncols
 
-let col t = List.Assoc.find t.cols ~equal:String.equal
+let get_col_by_name t = List.Assoc.find t.cols ~equal:String.equal
+let get_col t i =
+  List.nth t.cols i
+  |> Option.map ~f:snd
 
 let parse_header h =
   let labels = String.split ~on:'\t' h in
@@ -46,17 +49,17 @@ let rev_convert_col col =
   in
   let conv_opt f = conv (optionally f) in
   if List.mem col "NA" ~equal:String.equal then
-    try Maybe_int (conv_opt Int.of_string) with _ ->
-    try Maybe_float (conv_opt Float.of_string)
+    try Int_opts (conv_opt Int.of_string) with _ ->
+    try Float_opts (conv_opt Float.of_string)
     with _ ->
-      Maybe_string (
+      String_opts (
         List.map col ~f:Option.some
         |> Array.of_list
       )
   else
-    try Int (conv Int.of_string) with _ ->
-    try Float (conv Float.of_string) with _ ->
-    String (Array.of_list_rev col)
+    try Ints (conv Int.of_string) with _ ->
+    try Floats (conv Float.of_string) with _ ->
+    Strings (Array.of_list_rev col)
 
 let parse_lines ncols lines =
   let open Result.Monad_infix in
@@ -83,9 +86,9 @@ let%expect_test "Dataframe.parse_line ex1" =
   print_endline (show_parse_result got) ;
   [%expect {|
     (Ok (3,
-         [(Dataframe.String [|"a"; "a"; "c"|]);
-           (Dataframe.Float [|1.2; 1.2; -1.2|]);
-           (Dataframe.Maybe_int [|None; (Some 2); None|])])) |}]
+         [(Dataframe.Strings [|"a"; "a"; "c"|]);
+           (Dataframe.Floats [|1.2; 1.2; -1.2|]);
+           (Dataframe.Int_opts [|None; (Some 2); None|])])) |}]
 
 let from_file path =
   let open Result.Monad_infix in

lib/toolbox/dataframe.mli

 type t
 type column =
-  | Int of int array
-  | Maybe_int of int option array
-  | Float of float array
-  | Maybe_float of float option array
-  | String of string array
-  | Maybe_string of string option array
+  | Ints of int array
+  | Int_opts of int option array
+  | Floats of float array
+  | Float_opts of float option array
+  | Strings of string array
+  | String_opts of string option array
 
 val nrows : t -> int
 
 val ncols : t -> int
 
-val col : t -> string -> column option
+val get_col : t -> int -> column option
+
+val get_col_by_name : t -> string -> column option
 
 val from_file : string -> (t, [> `Msg of string]) result