Fix IP ranges labeling issue in pie chart (#17067) (#17144)

* Fix IP ranges labeling issue in pie chart

* Fix tests

src/ui/public/agg_response/hierarchical/__tests__/extract_buckets.js

@@ -38,7 +38,25 @@ describe('buildHierarchicalData()', function () {
       };
       const buckets = extractBuckets(bucket);
       expect(buckets).to.be.an(Array);
-      expect(buckets).to.be(bucket.buckets);
+      expect(buckets).to.eql(bucket.buckets);
+    });
+
+    it('should attach keys using agg.getKey for array of buckets', () => {
+      const bucket = {
+        buckets: [
+          { from: 10, doc_count: 1 },
+          { from: 20, doc_count: 2 }
+        ]
+      };
+      const agg = {
+        getKey(bucket) {
+          return bucket.from;
+        }
+      };
+      const buckets = extractBuckets(bucket, agg);
+      expect(buckets).to.have.length(2);
+      expect(buckets[0].key).to.be(10);
+      expect(buckets[1].key).to.be(20);
     });
 
   });

src/ui/public/agg_response/hierarchical/__tests__/transform_aggregation.js

@@ -20,7 +20,7 @@ describe('buildHierarchicalData()', function () {
           id: id,
           name: name,
           schema: { group: 'buckets' },
-          getKey: _.noop,
+          getKey: (bucket) => bucket.key,
           fieldFormatter: _.constant(String)
         };
       }

src/ui/public/agg_response/hierarchical/_extract_buckets.js

@@ -1,14 +1,22 @@
 import _ from 'lodash';
 
-export function extractBuckets(bucket, agg) {
-  if (bucket && _.isPlainObject(bucket.buckets)) {
-    return _.map(bucket.buckets, function (value, key) {
-      const item = _.cloneDeep(value);
-      item.key = agg ? agg.getKey(value, key) : key;
-      return item;
-    });
-
-  } else {
-    return bucket && bucket.buckets || [];
-  }
+function decorateWithKey(agg, bucket, key) {
+  const decorated = _.cloneDeep(bucket);
+  decorated.key = agg ? agg.getKey(bucket, key) : key;
+  return decorated;
+}
+
+export function extractBuckets(bucket, agg) {
+  if (bucket) {
+    if (_.isPlainObject(bucket.buckets)) {
+      return _.map(bucket.buckets, function (value, key) {
+        return decorateWithKey(agg, value, key);
+      });
+    } else if(_.isArray(bucket.buckets)) {
+      return bucket.buckets.map(value => {
+        return decorateWithKey(agg, value, value.key);
+      });
+    }
+  }
+  return bucket && bucket.buckets || [];
 }