diff --git a/src/hoc/treeTable/index.js b/src/hoc/treeTable/index.js
index 671e140..6711acd 100644
--- a/src/hoc/treeTable/index.js
+++ b/src/hoc/treeTable/index.js
@@ -49,8 +49,9 @@ export default Component => {
       const extra = {
         columns: columns.map(col => {
           let column = col
-          if (rest.pivotBy && rest.pivotBy.includes(col.accessor)) {
+          if (rest.pivotBy && (rest.pivotBy.includes(col.accessor) || rest.pivotBy.includes(col.id))) {
             column = {
+              id: col.id,
               accessor: col.accessor,
               width: `${treeTableIndent}px`,
               show: false,
diff --git a/src/index.js b/src/index.js
index 68e4c9c..1b24547 100644
--- a/src/index.js
+++ b/src/index.js
@@ -643,7 +643,7 @@ export default class ReactTable extends Methods(Lifecycle(Component)) {
                   key={`${i2}-${column.id}`}
                   className={classnames(
                     classes,
-                    !show && 'hidden',
+                    !cellInfo.expandable && !show && 'hidden',
                     cellInfo.expandable && 'rt-expandable',
                     (isBranch || isPreview) && 'rt-pivot'
                   )}