diff --git a/src/kibana/plugins/discover/components/field_chooser/field_chooser.js b/src/kibana/plugins/discover/components/field_chooser/field_chooser.js
index 4f625463db731c..144c856be9e0e2 100644
--- a/src/kibana/plugins/discover/components/field_chooser/field_chooser.js
+++ b/src/kibana/plugins/discover/components/field_chooser/field_chooser.js
@@ -146,7 +146,7 @@ define(function (require) {
 
           // Get all fields current in data set
           var currentFields = _.chain($scope.data).map(function (d) {
-            return _.keys(d.$$_flattened);
+            return _.keys($scope.indexPattern.flattenHit(d));
           }).flatten().unique().sort().value();
 
           _.each($scope.fields, function (field) {
diff --git a/test/unit/specs/apps/discover/directives/field_chooser.js b/test/unit/specs/apps/discover/directives/field_chooser.js
index 50b4b57b321ef2..ea7e82f87c611f 100644
--- a/test/unit/specs/apps/discover/directives/field_chooser.js
+++ b/test/unit/specs/apps/discover/directives/field_chooser.js
@@ -118,32 +118,35 @@ define(function (require) {
       });
 
       it('should not show the popular fields if there are not any', function (done) {
+
         // Re-init
         destroy();
+
+        _.each(indexPattern.fields, function (field) { field.count = 0;}); // Reset the popular fields
         init($elem, {
           columns: [],
           toggle: sinon.spy(),
           data: require('fixtures/hits'),
           filter: sinon.spy(),
-          indexPattern: indexPattern.fields
+          indexPattern: indexPattern
         });
 
         var section = getSections($elem);
-        // Remove the popular fields
+
         $scope.$digest();
         expect(section.popular.hasClass('ng-hide')).to.be(true);
         expect(section.popular.find('li:not(.sidebar-list-header)').length).to.be(0);
         done();
       });
 
-      it('should move the field into selected when setting field.display', function (done) {
+      it('should move the field into selected when it is added to the columns array', function (done) {
         var section = getSections($elem);
-        indexPattern.fields.byName.bytes.display = true;
+        $scope.columns.push('bytes');
         $scope.$digest();
         expect(section.selected.text()).to.contain('bytes');
         expect(section.popular.text()).to.not.contain('bytes');
 
-        indexPattern.fields.byName.ip.display = true;
+        $scope.columns.push('ip');
         $scope.$digest();
         expect(section.selected.text()).to.contain('ip\n');
         expect(section.unpopular.text()).to.not.contain('ip\n');