Implement slicing for categorical columns (#3379)

WIP for #1691
h2oai · May 2, 2023 · 390df8a · 390df8a
1 parent 7fd759c
commit 390df8a
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/src/core/column.h b/src/core/column.h
@@ -119,6 +119,7 @@ class Column
     size_t nrows() const noexcept;
     size_t na_count() const;
     const dt::Type& type() const noexcept;
+    const dt::Type& data_type() const noexcept;
     dt::LType ltype() const noexcept;
     dt::SType stype() const noexcept;
 

diff --git a/src/core/column/column_impl.cc b/src/core/column/column_impl.cc
@@ -240,7 +240,8 @@ const Type& ColumnImpl::data_type() const {
 }
 
 SType ColumnImpl::data_stype() const {
-  return data_type().stype();
+  return type_.is_categorical()? child(0).stype()
+                               : stype();
 }
 
 

diff --git a/src/core/column/column_impl.h b/src/core/column/column_impl.h
@@ -95,13 +95,10 @@ class ColumnImpl
     size_t nrows() const noexcept { return nrows_; }
     SType stype() const { return type_.stype(); }
     const Type& type() const { return type_; }
-
-    // For categorical columns these methods will return the stype/type
-    // of the data the column is backed up with. For all the other column types,
-    // these methods are equivalent to `stype()`/`type()`.
-    SType data_stype() const;
-    const Type& data_type() const;
-
+    const Type& data_type() const {
+      return type_.is_categorical()? child(0).type()
+                                   : type_;
+    }
     virtual bool is_virtual() const noexcept = 0;
     virtual bool computationally_expensive() const { return false; }
     virtual size_t memory_footprint() const noexcept = 0;