diff --git a/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.test.tsx b/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.test.tsx
index 5d02f0a2c759e..483446daed10f 100644
--- a/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.test.tsx
+++ b/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.test.tsx
@@ -24,7 +24,7 @@ import { IAggConfigs, IAggConfig } from 'src/plugins/data/public';
 import { DefaultEditorAggGroup, DefaultEditorAggGroupProps } from './agg_group';
 import { DefaultEditorAgg } from './agg';
 import { DefaultEditorAggAdd } from './agg_add';
-import { Schema } from '../schemas';
+import { ISchemas, Schemas } from '../schemas';
 import { EditorVisState } from './sidebar/state/reducers';
 
 jest.mock('@elastic/eui', () => ({
@@ -47,6 +47,7 @@ jest.mock('./agg_add', () => ({
 describe('DefaultEditorAgg component', () => {
   let defaultProps: DefaultEditorAggGroupProps;
   let aggs: IAggConfigs;
+  let schemas: ISchemas;
   let setTouched: jest.Mock;
   let setValidity: jest.Mock;
   let reorderAggs: jest.Mock;
@@ -55,6 +56,18 @@ describe('DefaultEditorAgg component', () => {
     setTouched = jest.fn();
     setValidity = jest.fn();
     reorderAggs = jest.fn();
+    schemas = new Schemas([
+      {
+        name: 'metrics',
+        group: 'metrics',
+        max: 1,
+      },
+      {
+        name: 'buckets',
+        group: 'buckets',
+        max: 1,
+      },
+    ]);
 
     aggs = {
       aggs: [
@@ -95,18 +108,7 @@ describe('DefaultEditorAgg component', () => {
       state: {
         data: { aggs },
       } as EditorVisState,
-      schemas: [
-        {
-          name: 'metrics',
-          group: 'metrics',
-          max: 1,
-        } as Schema,
-        {
-          name: 'buckets',
-          group: 'buckets',
-          max: 1,
-        } as Schema,
-      ],
+      schemas: schemas.metrics,
       setTouched,
       setValidity,
       reorderAggs,
@@ -133,6 +135,7 @@ describe('DefaultEditorAgg component', () => {
 
   it('should last bucket has truthy isLastBucket prop', () => {
     defaultProps.groupName = 'buckets';
+    defaultProps.schemas = schemas.buckets;
     const comp = mount(<DefaultEditorAggGroup {...defaultProps} />);
     const lastAgg = comp.find(DefaultEditorAgg).last();
 
@@ -154,6 +157,8 @@ describe('DefaultEditorAgg component', () => {
 
   it('should show add button when schemas count is less than max', () => {
     defaultProps.groupName = 'buckets';
+    defaultProps.schemas = schemas.buckets;
+    defaultProps.schemas[0].max = 2;
     const comp = shallow(<DefaultEditorAggGroup {...defaultProps} />);
 
     expect(comp.find(DefaultEditorAggAdd).exists()).toBeTruthy();
diff --git a/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.tsx b/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.tsx
index f50abc3ebb599..792595fd421f6 100644
--- a/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.tsx
+++ b/src/legacy/core_plugins/vis_default_editor/public/components/agg_group.tsx
@@ -41,7 +41,7 @@ import {
   getEnabledMetricAggsCount,
 } from './agg_group_helper';
 import { aggGroupReducer, initAggsState, AGGS_ACTION_KEYS } from './agg_group_state';
-import { Schema, getSchemasByGroup } from '../schemas';
+import { Schema } from '../schemas';
 import { TimeRange } from '../../../../../plugins/data/public';
 
 export interface DefaultEditorAggGroupProps extends DefaultEditorAggCommonProps {
@@ -73,7 +73,7 @@ function DefaultEditorAggGroup({
 }: DefaultEditorAggGroupProps) {
   const groupNameLabel = (search.aggs.aggGroupNamesMap() as any)[groupName];
   // e.g. buckets can have no aggs
-  const schemaNames = getSchemasByGroup(schemas, groupName).map(s => s.name);
+  const schemaNames = schemas.map(s => s.name);
   const group: IAggConfig[] = useMemo(
     () =>
       state.data.aggs!.aggs.filter(
diff --git a/src/legacy/core_plugins/vis_default_editor/public/components/sidebar/sidebar.tsx b/src/legacy/core_plugins/vis_default_editor/public/components/sidebar/sidebar.tsx
index 071e10682e22c..b24486a12fd24 100644
--- a/src/legacy/core_plugins/vis_default_editor/public/components/sidebar/sidebar.tsx
+++ b/src/legacy/core_plugins/vis_default_editor/public/components/sidebar/sidebar.tsx
@@ -31,8 +31,7 @@ import { DefaultEditorAggCommonProps } from '../agg_common_props';
 import { SidebarTitle } from './sidebar_title';
 import { PersistedState } from '../../../../../../plugins/visualizations/public';
 import { SavedSearch } from '../../../../../../plugins/discover/public';
-import { AggGroupNames } from '../../../../../../plugins/data/public';
-import { getSchemasByGroup } from '../../schemas';
+import { Schema } from '../../schemas';
 import { TimeRange } from '../../../../../../plugins/data/public';
 
 interface DefaultEditorSideBarProps {
@@ -66,9 +65,7 @@ function DefaultEditorSideBar({
   const responseAggs = useMemo(() => (state.data.aggs ? state.data.aggs.getResponseAggs() : []), [
     state.data.aggs,
   ]);
-  const metricSchemas = getSchemasByGroup(vis.type.schemas.all || [], AggGroupNames.Metrics).map(
-    s => s.name
-  );
+  const metricSchemas = (vis.type.schemas.metrics || []).map((s: Schema) => s.name);
   const metricAggs = useMemo(
     () => responseAggs.filter(agg => metricSchemas.includes(get(agg, 'schema'))),
     [responseAggs, metricSchemas]
diff --git a/src/legacy/core_plugins/vis_default_editor/public/schemas.ts b/src/legacy/core_plugins/vis_default_editor/public/schemas.ts
index 94e3ad6023f4e..4e632da44afc0 100644
--- a/src/legacy/core_plugins/vis_default_editor/public/schemas.ts
+++ b/src/legacy/core_plugins/vis_default_editor/public/schemas.ts
@@ -17,11 +17,10 @@
  * under the License.
  */
 
-import _ from 'lodash';
+import _, { defaults } from 'lodash';
 
 import { Optional } from '@kbn/utility-types';
 
-import { IndexedArray } from 'ui/indexed_array';
 import { AggGroupNames, AggParam, IAggGroupNames } from '../../../../plugins/data/public';
 
 export interface ISchemas {
@@ -45,9 +44,10 @@ export interface Schema {
   aggSettings?: any;
 }
 
-export class Schemas {
-  // @ts-ignore
-  all: IndexedArray<Schema>;
+export class Schemas implements ISchemas {
+  all: Schema[] = [];
+  [AggGroupNames.Buckets]: Schema[] = [];
+  [AggGroupNames.Metrics]: Schema[] = [];
 
   constructor(
     schemas: Array<
@@ -70,7 +70,7 @@ export class Schemas {
           ] as AggParam[];
         }
 
-        _.defaults(schema, {
+        defaults(schema, {
           min: 0,
           max: Infinity,
           group: AggGroupNames.Buckets,
@@ -83,22 +83,12 @@ export class Schemas {
         return schema as Schema;
       })
       .tap((fullSchemas: Schema[]) => {
-        this.all = new IndexedArray({
-          index: ['name'],
-          group: ['group'],
-          immutable: true,
-          initialSet: fullSchemas,
-        });
+        this.all = fullSchemas;
       })
       .groupBy('group')
       .forOwn((group, groupName) => {
         // @ts-ignore
-        this[groupName] = new IndexedArray({
-          index: ['name'],
-          immutable: true,
-          // @ts-ignore
-          initialSet: group,
-        });
+        this[groupName] = group;
       })
       .commit();
   }
@@ -107,7 +97,3 @@ export class Schemas {
 export const getSchemaByName = (schemas: Schema[], schemaName?: string) => {
   return schemas.find(s => s.name === schemaName) || ({} as Schema);
 };
-
-export const getSchemasByGroup = (schemas: Schema[], schemaGroup?: string) => {
-  return schemas.filter(s => s.group === schemaGroup);
-};