Allow using column selector on cat columns (#3240)

WIP for #1691
h2oai · Feb 22, 2022 · 6376ea0 · 6376ea0
1 parent 06d3f16
commit 6376ea0
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/src/core/expr/eval_context.cc b/src/core/expr/eval_context.cc
@@ -498,7 +498,9 @@ static void _vivify_workframe(const Workframe& wf) {
   if (wf.nrows() == 0) return;
   for (size_t i = 0; i < wf.ncols(); ++i) {
     const Column& col = wf.get_column(i);
-    switch (col.stype()) {
+    dt::SType st = col.type().is_categorical()? col.child(0).stype()
+                                              : col.stype();
+    switch (st) {
       case SType::VOID:
       case SType::BOOL:
       case SType::INT8:    _vivify_column<int8_t>(col); break;

diff --git a/tests/types/test-categorical.py b/tests/types/test-categorical.py
@@ -139,6 +139,7 @@ def test_create_from_void(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -152,6 +153,7 @@ def test_create_from_bools(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -165,6 +167,7 @@ def test_create_from_ints(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 def test_create_from_ints_large():
@@ -174,6 +177,7 @@ def test_create_from_ints_large():
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -187,6 +191,7 @@ def test_create_from_floats(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -200,6 +205,7 @@ def test_create_from_strings(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -214,6 +220,7 @@ def test_create_from_dates(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 @pytest.mark.parametrize('t', [dt.Type.cat8,
@@ -230,6 +237,23 @@ def test_create_from_times(t):
     assert DT1.shape == DT2.shape
     assert DT1.names == DT2.names
     assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
+
+
+@pytest.mark.parametrize('t', [dt.Type.cat8,
+                               dt.Type.cat16,
+                               dt.Type.cat32])
+def test_create_multicolumn(t):
+    src = [[6, 8, 8, None, 20],
+           ["six", "eight", "eight", None, "twenty"]]
+    DT1 = dt.Frame(src)
+    DT2 = dt.Frame(src, types = [t(dt.Type.int32), t(dt.Type.str32)])
+    assert DT2[:, 0].type == t(dt.Type.int32)
+    assert DT2[:, 1].type == t(dt.Type.str32)
+    assert DT1.shape == DT2.shape
+    assert DT1.names == DT2.names
+    assert DT1.to_list() == DT2.to_list()
+    assert_equals(DT2, DT2[:, :])
 
 
 #-------------------------------------------------------------------------------