Merge pull request modin-project#50 from intel-go/ienkovich/cat-codes…

…-null fix cat.codes for NULLs
dchigarev · Jun 23, 2020 · 0084e8d · 0084e8d
2 parents 54c5df5 + f85e682
commit 0084e8d
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/modin/experimental/engines/omnisci_on_ray/frame/data.py b/modin/experimental/engines/omnisci_on_ray/frame/data.py
@@ -586,7 +586,12 @@ def cat_codes(self):
 
         col = self.columns[-1]
         exprs = self._index_exprs()
-        exprs[col] = OpExpr("KEY_FOR_STRING", [self.ref(col)], _get_dtype("int32"))
+        col_expr = self.ref(col)
+        code_expr = OpExpr("KEY_FOR_STRING", [col_expr], _get_dtype("int32"))
+        null_val = LiteralExpr(np.int32(-1))
+        exprs[col] = build_if_then_else(
+            col_expr.is_null(), null_val, code_expr, _get_dtype("int32")
+        )
 
         return self.__constructor__(
             columns=self.columns,

diff --git a/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py b/modin/experimental/engines/omnisci_on_ray/test/test_dataframe.py
@@ -871,7 +871,7 @@ def test_dt_month(self):
 
 class TestCategory:
     data = {
-        "a": ["str1", "str2", "str1", "str3", "str2"],
+        "a": ["str1", "str2", "str1", "str3", "str2", None],
     }
 
     def test_cat_codes(self):