Add support for categorical types (#121)

* add support for categorical types

* fix some clippy issues

* fix issue with guide

guide/expressions.ml

@@ -561,11 +561,11 @@ let%expect_test "Casting" =
 let%expect_test "Aggregation" =
   let schema =
     Schema.create
-      [ "first_name", Utf8
-      ; "gender", Utf8
-      ; "type", Utf8
-      ; "state", Utf8
-      ; "party", Utf8
+      [ "first_name", Categorical None
+      ; "gender", Categorical None
+      ; "type", Categorical None
+      ; "state", Categorical None
+      ; "party", Categorical None
       ; "birthday", Date
       ]
   in
@@ -582,7 +582,7 @@ let%expect_test "Aggregation" =
     ┌────────────┬────────────┬────────────┬────────┬───┬───────────┬───────────┬──────────┬───────────┐
     │ last_name  ┆ first_name ┆ middle_nam ┆ suffix ┆ … ┆ ballotped ┆ washingto ┆ icpsr_id ┆ wikipedia │
     │ ---        ┆ ---        ┆ e          ┆ ---    ┆   ┆ ia_id     ┆ n_post_id ┆ ---      ┆ _id       │
-    │ str        ┆ str        ┆ ---        ┆ str    ┆   ┆ ---       ┆ ---       ┆ i64      ┆ ---       │
+    │ str        ┆ cat        ┆ ---        ┆ str    ┆   ┆ ---       ┆ ---       ┆ i64      ┆ ---       │
     │            ┆            ┆ str        ┆        ┆   ┆ str       ┆ str       ┆          ┆ str       │
     ╞════════════╪════════════╪════════════╪════════╪═══╪═══════════╪═══════════╪══════════╪═══════════╡
     │ Bassett    ┆ Richard    ┆ null       ┆ null   ┆ … ┆ null      ┆ null      ┆ 507      ┆ Richard   │
@@ -627,7 +627,7 @@ let%expect_test "Aggregation" =
     ┌────────────┬───────┬───────────────────┬───────────┐
     │ first_name ┆ count ┆ gender            ┆ last_name │
     │ ---        ┆ ---   ┆ ---               ┆ ---       │
-    │ str        ┆ u32   ┆ list[str]         ┆ str       │
+    │ cat        ┆ u32   ┆ list[cat]         ┆ str       │
     ╞════════════╪═══════╪═══════════════════╪═══════════╡
     │ John       ┆ 1256  ┆ ["M", "M", … "M"] ┆ Walker    │
     │ William    ┆ 1022  ┆ ["M", "M", … "M"] ┆ Few       │
@@ -655,7 +655,7 @@ let%expect_test "Aggregation" =
     ┌───────┬──────┬─────┐
     │ state ┆ anti ┆ pro │
     │ ---   ┆ ---  ┆ --- │
-    │ str   ┆ u32  ┆ u32 │
+    │ cat   ┆ u32  ┆ u32 │
     ╞═══════╪══════╪═════╡
     │ NJ    ┆ 0    ┆ 3   │
     │ CT    ┆ 0    ┆ 3   │
@@ -684,7 +684,7 @@ let%expect_test "Aggregation" =
     ┌───────┬─────────────────────┬───────┐
     │ state ┆ party               ┆ count │
     │ ---   ┆ ---                 ┆ ---   │
-    │ str   ┆ str                 ┆ u32   │
+    │ cat   ┆ cat                 ┆ u32   │
     ╞═══════╪═════════════════════╪═══════╡
     │ NJ    ┆ Pro-Administration  ┆ 3     │
     │ VA    ┆ Anti-Administration ┆ 3     │
@@ -720,7 +720,7 @@ let%expect_test "Aggregation" =
     ┌───────┬────────────────┬────────────────┬────────┬──────────┐
     │ state ┆ avg M birthday ┆ avg F birthday ┆ # male ┆ # female │
     │ ---   ┆ ---            ┆ ---            ┆ ---    ┆ ---      │
-    │ str   ┆ f64            ┆ f64            ┆ u32    ┆ u32      │
+    │ cat   ┆ f64            ┆ f64            ┆ u32    ┆ u32      │
     ╞═══════╪════════════════╪════════════════╪════════╪══════════╡
     │ DE    ┆ 182.593407     ┆ null           ┆ 97     ┆ 0        │
     │ VA    ┆ 192.542781     ┆ 66.2           ┆ 430    ┆ 5        │
@@ -749,7 +749,7 @@ let%expect_test "Aggregation" =
     ┌───────┬──────────────────┬───────────────────────┐
     │ state ┆ youngest         ┆ oldest                │
     │ ---   ┆ ---              ┆ ---                   │
-    │ str   ┆ str              ┆ str                   │
+    │ cat   ┆ str              ┆ str                   │
     ╞═══════╪══════════════════╪═══════════════════════╡
     │ NC    ┆ Madison Cawthorn ┆ John Ashe             │
     │ IA    ┆ Abby Finkenauer  ┆ Bernhart Henn         │
@@ -778,7 +778,7 @@ let%expect_test "Aggregation" =
     ┌───────┬──────────────────┬───────────────────────┬────────────────────┐
     │ state ┆ youngest         ┆ oldest                ┆ alphabetical_first │
     │ ---   ┆ ---              ┆ ---                   ┆ ---                │
-    │ str   ┆ str              ┆ str                   ┆ str                │
+    │ cat   ┆ str              ┆ str                   ┆ str                │
     ╞═══════╪══════════════════╪═══════════════════════╪════════════════════╡
     │ NC    ┆ Madison Cawthorn ┆ John Ashe             ┆ Abraham Rencher    │
     │ IA    ┆ Abby Finkenauer  ┆ Bernhart Henn         ┆ Abby Finkenauer    │
@@ -797,29 +797,35 @@ let%expect_test "Aggregation" =
              ; get_person |> last |> alias ~name:"oldest"
              ; get_person |> sort |> first |> alias ~name:"alphabetical_first"
              ; col "gender"
-               |> sort_by ~by:[ col "first_name" ]
+               |> sort_by
+                    ~by:
+                      [ (* The guide uses "first_name" to sort by, but I'm guessing
+                           there's an nondeterminism bug causing output to be unstable
+                           if we have multiple sorts or something, so I suspect
+                           [get_person] is what we actually want *)
+                        get_person
+                      ]
                |> first
                |> alias ~name:"gender"
              ]
-    |> Lazy_frame.sort ~by_column:"state"
     |> Lazy_frame.limit ~n:5
     |> Lazy_frame.collect_exn
   in
   Data_frame.print df;
   [%expect
     {|
     shape: (5, 5)
-    ┌───────┬──────────────────┬────────────────┬────────────────────┬────────┐
-    │ state ┆ youngest         ┆ oldest         ┆ alphabetical_first ┆ gender │
-    │ ---   ┆ ---              ┆ ---            ┆ ---                ┆ ---    │
-    │ str   ┆ str              ┆ str            ┆ str                ┆ str    │
-    ╞═══════╪══════════════════╪════════════════╪════════════════════╪════════╡
-    │ AK    ┆ Mark Begich      ┆ Thomas Cale    ┆ Anthony Dimond     ┆ M      │
-    │ AL    ┆ Martha Roby      ┆ John McKee     ┆ Albert Goodwyn     ┆ M      │
-    │ AR    ┆ Tim Griffin      ┆ Archibald Yell ┆ Albert Rust        ┆ M      │
-    │ AS    ┆ Eni Faleomavaega ┆ Fofó Sunia     ┆ Eni Faleomavaega   ┆ M      │
-    │ AZ    ┆ Ben Quayle       ┆ Coles Bashford ┆ Ann Kirkpatrick    ┆ F      │
-    └───────┴──────────────────┴────────────────┴────────────────────┴────────┘ |}]
+    ┌───────┬──────────────────┬───────────────────────┬────────────────────┬────────┐
+    │ state ┆ youngest         ┆ oldest                ┆ alphabetical_first ┆ gender │
+    │ ---   ┆ ---              ┆ ---                   ┆ ---                ┆ ---    │
+    │ cat   ┆ str              ┆ str                   ┆ str                ┆ cat    │
+    ╞═══════╪══════════════════╪═══════════════════════╪════════════════════╪════════╡
+    │ NC    ┆ Madison Cawthorn ┆ John Ashe             ┆ Abraham Rencher    ┆ M      │
+    │ IA    ┆ Abby Finkenauer  ┆ Bernhart Henn         ┆ Abby Finkenauer    ┆ F      │
+    │ MI    ┆ Peter Meijer     ┆ Edward Bradley        ┆ Aaron Bliss        ┆ M      │
+    │ CA    ┆ Katie Hill       ┆ Edward Gilbert        ┆ Aaron Sargent      ┆ M      │
+    │ NY    ┆ Mondaire Jones   ┆ Cornelius Schoonmaker ┆ A. Foster          ┆ M      │
+    └───────┴──────────────────┴───────────────────────┴────────────────────┴────────┘ |}]
 ;;
 
 (* Examples from https://pola-rs.github.io/polars-book/user-guide/expressions/null/ *)

lib/data_type.ml

@@ -1,40 +1,38 @@
 open! Core
 
-module T = struct
-  type t =
-    | Boolean
-    | UInt8
-    | UInt16
-    | UInt32
-    | UInt64
-    | Int8
-    | Int16
-    | Int32
-    | Int64
-    | Float32
-    | Float64
-    | Utf8
-    | Binary
-    | Date
-    | Datetime of Time_unit.t * Tz.t option
-    | Duration of Time_unit.t
-    | Time
-    | List of t
-    (* We want this branch to be tested very well, since code dealing with
-       this recursive case is usually the most non-trivial portion of the
-       logic. *)
-    [@quickcheck.weight 10.]
-    | Null
-    | Struct of (string * t) list
-    | Unknown
-  [@@deriving compare, sexp, quickcheck]
-end
-
-include T
-include Sexpable.To_stringable (T)
+type t =
+  | Boolean
+  | UInt8
+  | UInt16
+  | UInt32
+  | UInt64
+  | Int8
+  | Int16
+  | Int32
+  | Int64
+  | Float32
+  | Float64
+  | Utf8
+  | Binary
+  | Date
+  | Datetime of Time_unit.t * Tz.t option
+  | Duration of Time_unit.t
+  | Time
+  | List of t
+  (* We want this branch to be tested very well, since code dealing with
+     this recursive case is usually the most non-trivial portion of the
+     logic. *)
+  [@quickcheck.weight 10.]
+  | Null
+  | Categorical of (Rev_mapping.t[@compare.ignore]) option [@quickcheck.do_not_generate]
+  | Struct of (string * t) list
+  | Unknown
+[@@deriving compare, sexp_of, quickcheck]
+
+let to_string t = Sexp.to_string ([%sexp_of: t] t)
 
 module Typed = struct
-  type untyped = t [@@deriving compare, sexp, quickcheck]
+  type untyped = t [@@deriving compare, sexp_of, quickcheck]
 
   (* TODO: Consider mapping to smaller OCaml values like Int8, Float32, etc instead of
      casting up *)
@@ -157,7 +155,7 @@ module Typed = struct
     | Duration time_unit -> Some (T (Duration time_unit))
     | Time -> Some (T Time)
     | List t -> of_untyped t |> Option.map ~f:(fun (T t) -> T (List t))
-    | Null | Struct _ | Unknown -> None
+    | Null | Categorical _ | Struct _ | Unknown -> None
   ;;
 
   let rec sexp_of_packed (T t) =

lib/data_type.mli

@@ -20,11 +20,12 @@ type t =
   | Time
   | List of t
   | Null
+  | Categorical of Rev_mapping.t option
   | Struct of (string * t) list
   | Unknown
-[@@deriving compare, sexp, quickcheck]
+[@@deriving compare, sexp_of, quickcheck]
 
-include Stringable.S with type t := t
+val to_string : t -> string
 
 module Typed : sig
     type untyped

lib/polars.ml

@@ -14,5 +14,6 @@ module Naive_time = Naive_time
 module Schema = Schema
 module Series = Series
 module Sql_context = Sql_context
+module Rev_mapping = Rev_mapping
 module Time_unit = Time_unit
 module Tz = Tz

lib/rev_mapping.ml

@@ -0,0 +1,19 @@
+open! Core
+
+type t
+
+let quickcheck_shrinker = Base_quickcheck.Shrinker.atomic
+
+let quickcheck_observer =
+  Base_quickcheck.Observer.of_hash_fold (fun hash_state _t -> hash_state)
+;;
+
+external get_categories : t -> string list = "rust_rev_mapping_get_categories"
+
+let sexp_of_t t =
+  get_categories t
+  |> (* It's difficult to get tests to be deterministic wrt the ordering of
+        categories, so we force them to be sorted here as a hacky workaround. *)
+  List.sort ~compare:String.compare
+  |> [%sexp_of: string list]
+;;

lib/rev_mapping.mli

@@ -0,0 +1,7 @@
+open! Core
+
+type t [@@deriving sexp_of]
+
+val quickcheck_shrinker : t Quickcheck.Shrinker.t
+val quickcheck_observer : t Quickcheck.Observer.t
+val get_categories : t -> string list

lib/schema.ml

@@ -6,4 +6,3 @@ external create : (string * Data_type.t) list -> t = "rust_schema_create"
 external to_fields : t -> (string * Data_type.t) list = "rust_schema_to_fields"
 
 let sexp_of_t t = to_fields t |> [%sexp_of: (string * Data_type.t) list]
-let t_of_sexp sexp = [%of_sexp: (string * Data_type.t) list] sexp |> create

lib/schema.mli

@@ -1,6 +1,6 @@
 open! Core
 
-type t [@@deriving sexp]
+type t [@@deriving sexp_of]
 
 val create : (string * Data_type.t) list -> t
 val to_fields : t -> (string * Data_type.t) list

lib/series.ml

@@ -211,6 +211,15 @@ module T = struct
     map input_data_type output_data_type t ~f |> Result.ok_exn
   ;;
 
+  external cast
+    :  t
+    -> to_:Data_type.t
+    -> strict:bool
+    -> (t, string) result
+    = "rust_series_cast"
+
+  let cast ?(strict = true) t ~to_ = cast t ~to_ ~strict |> Utils.string_result_ok_exn
+
   external name : t -> string = "rust_series_name"
   external rename : t -> name:string -> unit = "rust_series_rename"
   external dtype : t -> Data_type.t = "rust_series_dtype"

lib/series.mli

@@ -75,6 +75,7 @@ val map
   -> f:('a option -> 'b option)
   -> t
 
+val cast : ?strict:bool -> t -> to_:Data_type.t -> t
 val name : t -> string
 val rename : t -> name:string -> unit
 val dtype : t -> Data_type.t

rust/polars-ocaml/Cargo.toml

@@ -35,6 +35,7 @@ features = [
     "dtype-i8",
     "dtype-u16",
     "dtype-u8",
+    "dtype-categorical",
     "dynamic_groupby",
     "horizontal_concat",
     "interpolate",

rust/polars-ocaml/src/expr.rs

@@ -5,7 +5,6 @@ use chrono::{Duration, NaiveDate, NaiveDateTime, NaiveTime};
 use ocaml_interop::{
     DynBox, OCaml, OCamlBytes, OCamlFloat, OCamlInt, OCamlList, OCamlRef, OCamlRuntime, ToOCaml,
 };
-use polars::lazy::dsl::GetOutput;
 use polars::prelude::*;
 use polars::series::IsSorted;
 use polars_ocaml_macros::ocaml_interop_export;

rust/polars-ocaml/src/misc.rs

@@ -31,6 +31,17 @@ fn rust_schema_to_fields(
     fields.to_ocaml(cr)
 }
 
+#[ocaml_interop_export]
+fn rust_rev_mapping_get_categories(
+    cr: &mut &mut OCamlRuntime,
+    rev_mapping: OCamlRef<DynBox<Arc<RevMapping>>>,
+) -> OCaml<OCamlList<String>> {
+    let Abstract(rev_mapping) = rev_mapping.to_rust(cr);
+    let rev_mapping: Vec<_> = rev_mapping.get_categories().values_iter().collect();
+
+    rev_mapping.to_ocaml(cr)
+}
+
 #[ocaml_interop_export]
 fn rust_test_panic(cr: &mut &mut OCamlRuntime, error_message: OCamlRef<String>) -> OCaml<()> {
     let error_message: String = error_message.to_rust(cr);

rust/polars-ocaml/src/polars_types.rs

@@ -5,8 +5,8 @@ use ocaml_interop::{
     polymorphic_variant_tag_hash, DynBox, FromOCaml, OCaml, OCamlInt, OCamlList, OCamlRuntime,
     ToOCaml,
 };
+use polars::prelude::*;
 use polars::series::IsSorted;
-use polars::{lazy::dsl::WindowMapping, prelude::*};
 use smartstring::{LazyCompact, SmartString};
 
 #[derive(Debug, Clone)]
@@ -71,6 +71,10 @@ unsafe impl FromOCaml<DataType> for PolarsDataType {
                     DataType::List(Box::new(datatype))
                 },
                 DataType::Null,
+                DataType::Categorical(local_rev_mapping_opt: Option<DynBox<Arc<RevMapping>>>) => {
+                    let local_rev_mapping_opt: Option<Abstract<Arc<RevMapping>>> = local_rev_mapping_opt;
+                    DataType::Categorical(local_rev_mapping_opt.map(Abstract::get))
+                },
                 DataType::Struct(fields: OCamlList<(String, DataType)>) => {
                     let fields_: Vec<(String, PolarsDataType)> = fields;
                     let fields: Vec<Field> =
@@ -134,12 +138,16 @@ unsafe impl ToOCaml<DataType> for PolarsDataType {
                     ocaml_alloc_tagged_block!(cr, 2, datatype: DataType)
                 }
                 DataType::Null => ocaml_value(cr, 15),
+                DataType::Categorical(local_rev_mapping_opt) => {
+                    let local_rev_mapping_opt = local_rev_mapping_opt.clone().map(Abstract);
+                    ocaml_alloc_tagged_block!(cr, 3, local_rev_mapping_opt: Option<DynBox<Arc<RevMapping>>>)
+                }
                 DataType::Struct(fields) => {
                     let fields: Vec<(String, PolarsDataType)> = fields
                         .iter()
                         .map(|field| (field.name.to_string(), PolarsDataType(field.dtype.clone())))
                         .collect();
-                    ocaml_alloc_tagged_block!(cr, 3, fields: OCamlList<(String, DataType)>)
+                    ocaml_alloc_tagged_block!(cr, 4, fields: OCamlList<(String, DataType)>)
                 }
                 DataType::Unknown => ocaml_value(cr, 16),
             }