From 2763a4656b4ccfd86677b31c85149dff2d70d27f Mon Sep 17 00:00:00 2001
From: Ryan Curtin <ryan@ratml.org>
Date: Mon, 7 Jun 2021 16:34:30 -0400
Subject: [PATCH] Compute number of categories in categorical data.

---
 src/mlpack/bindings/julia/julia_util.cpp | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/mlpack/bindings/julia/julia_util.cpp b/src/mlpack/bindings/julia/julia_util.cpp
index ac8663eab0..e502527824 100644
--- a/src/mlpack/bindings/julia/julia_util.cpp
+++ b/src/mlpack/bindings/julia/julia_util.cpp
@@ -193,13 +193,37 @@ void IO_SetParamMatWithInfo(const char* paramName,
                             const bool pointsAreRows)
 {
   data::DatasetInfo d(pointsAreRows ? cols : rows);
+  bool hasCategoricals = false;
   for (size_t i = 0; i < d.Dimensionality(); ++i)
   {
     d.Type(i) = (dimensions[i]) ? data::Datatype::categorical :
         data::Datatype::numeric;
+    if (dimensions[i])
+      hasCategoricals = true;
   }
 
   arma::mat m(memptr, arma::uword(rows), arma::uword(cols), false, true);
+
+  // Do we need to find how many categories we have?
+  if (hasCategoricals)
+  {
+    arma::vec maxs = arma::max(m, 1);
+
+    for (size_t i = 0; i < d.Dimensionality(); ++i)
+    {
+      if (dimensions[i])
+      {
+        // Map the right number of objects.
+        for (size_t j = 0; j < (size_t) maxs[i]; ++j)
+        {
+          std::ostringstream oss;
+          oss << j;
+          d.MapString<double>(oss.str(), i);
+        }
+      }
+    }
+  }
+
   std::get<0>(IO::GetParam<std::tuple<data::DatasetInfo, arma::mat>>(
       paramName)) = std::move(d);
   std::get<1>(IO::GetParam<std::tuple<data::DatasetInfo, arma::mat>>(