[VisBuilder] Add Capability to generate dynamic vega

In this PR, we add the capability for Visbuilder to generate dynamic Vega and Vega-Lite
specifications based on user settings and aggregation configurations.

* developed functions buildVegaSpecViaVega and buildVegaSpecViaVegaLite
that can create either Vega or Vega-Lite specifications depending on the complexity
of the visualization.
* added VegaSpec and VegaLiteSpec interfaces to provide better type checking
* broken down the specification building into smaller, reusable components
(like buildEncoding, buildMark, buildLegend, buildTooltip) to make the code
more maintainable and easier to extend.
* added flattenDataHandler to prepare and transform data for use in Vega visualizations

Issue Resolve
#7067

Signed-off-by: Anan Zhuang <[email protected]>

src/plugins/advanced_settings/public/management_app/lib/get_category_name.ts

@@ -48,6 +48,9 @@ const names: Record<string, string> = {
   visualizations: i18n.translate('advancedSettings.categoryNames.visualizationsLabel', {
     defaultMessage: 'Visualizations',
   }),
+  visbuilder: i18n.translate('advancedSettings.categoryNames.visbuilderLabel', {
+    defaultMessage: 'VisBuilder',
+  }),
   discover: i18n.translate('advancedSettings.categoryNames.discoverLabel', {
     defaultMessage: 'Discover',
   }),

src/plugins/expressions/public/index.ts

@@ -134,3 +134,4 @@ export {
   UnmappedTypeStrings,
   ExpressionValueRender as Render,
 } from '../common';
+export { getExpressionsService } from './services';

src/plugins/vis_builder/common/constants.ts

@@ -0,0 +1,6 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export const VISBUILDER_ENABLE_VEGA_SETTING = 'visbuilder:enableVega';

src/plugins/vis_builder/public/application/components/workspace.tsx

@@ -13,6 +13,7 @@ import { validateSchemaState, validateAggregations } from '../utils/validations'
 import { useTypedDispatch, useTypedSelector, setUIStateState } from '../utils/state_management';
 import { useAggs, useVisualizationType } from '../utils/use';
 import { PersistedState } from '../../../../visualizations/public';
+import { VISBUILDER_ENABLE_VEGA_SETTING } from '../../../common/constants';
 
 import hand_field from '../../assets/hand_field.svg';
 import fields_bg from '../../assets/fields_bg.svg';
@@ -27,6 +28,7 @@ export const WorkspaceUI = () => {
       notifications: { toasts },
       data,
       uiActions,
+      uiSettings,
     },
   } = useOpenSearchDashboards<VisBuilderServices>();
   const { toExpression, ui } = useVisualizationType();
@@ -37,6 +39,7 @@ export const WorkspaceUI = () => {
     filters: data.query.filterManager.getFilters(),
     timeRange: data.query.timefilter.timefilter.getTime(),
   });
+  const useVega = uiSettings.get(VISBUILDER_ENABLE_VEGA_SETTING);
   const rootState = useTypedSelector((state) => state);
   const dispatch = useTypedDispatch();
   // Visualizations require the uiState object to persist even when the expression changes
@@ -81,12 +84,20 @@ export const WorkspaceUI = () => {
         return;
       }
 
-      const exp = await toExpression(rootState, searchContext);
+      const exp = await toExpression(rootState, searchContext, useVega);
       setExpression(exp);
     }
 
     loadExpression();
-  }, [rootState, toExpression, toasts, ui.containerConfig.data.schemas, searchContext, aggConfigs]);
+  }, [
+    rootState,
+    toExpression,
+    toasts,
+    ui.containerConfig.data.schemas,
+    searchContext,
+    aggConfigs,
+    useVega,
+  ]);
 
   useLayoutEffect(() => {
     const subscription = data.query.state$.subscribe(({ state }) => {

src/plugins/vis_builder/public/plugin.ts

@@ -56,6 +56,7 @@ import {
   withNotifyOnErrors,
 } from '../../opensearch_dashboards_utils/public';
 import { opensearchFilters } from '../../data/public';
+import { createRawDataVisFn } from './visualizations/vega/utils/expression_helper';
 
 export class VisBuilderPlugin
   implements
@@ -74,7 +75,7 @@ export class VisBuilderPlugin
 
   public setup(
     core: CoreSetup<VisBuilderPluginStartDependencies, VisBuilderStart>,
-    { embeddable, visualizations, data }: VisBuilderPluginSetupDependencies
+    { embeddable, visualizations, data, expressions: exp }: VisBuilderPluginSetupDependencies
   ) {
     const { appMounted, appUnMounted, stop: stopUrlTracker } = createOsdUrlTracker({
       baseUrl: core.http.basePath.prepend(`/app/${PLUGIN_ID}`),
@@ -107,6 +108,7 @@ export class VisBuilderPlugin
     // Register Default Visualizations
     const typeService = this.typeService;
     registerDefaultTypes(typeService.setup());
+    exp.registerFunction(createRawDataVisFn());
 
     // Register the plugin to core
     core.application.register({

src/plugins/vis_builder/public/services/type_service/types.ts

@@ -31,6 +31,7 @@ export interface VisualizationTypeOptions<T = any> {
   };
   readonly toExpression: (
     state: RenderState,
-    searchContext: IExpressionLoaderParams['searchContext']
+    searchContext: IExpressionLoaderParams['searchContext'],
+    useVega: boolean
   ) => Promise<string | undefined>;
 }

src/plugins/vis_builder/public/services/type_service/visualization_type.tsx

@@ -18,7 +18,8 @@ export class VisualizationType implements IVisualizationType {
   public readonly ui: IVisualizationType['ui'];
   public readonly toExpression: (
     state: RenderState,
-    searchContext: IExpressionLoaderParams['searchContext']
+    searchContext: IExpressionLoaderParams['searchContext'],
+    useVega: boolean
   ) => Promise<string | undefined>;
 
   constructor(options: VisualizationTypeOptions) {

src/plugins/vis_builder/public/types.ts

@@ -8,7 +8,7 @@ import { SavedObject, SavedObjectsStart } from '../../saved_objects/public';
 import { EmbeddableSetup, EmbeddableStart } from '../../embeddable/public';
 import { DashboardStart } from '../../dashboard/public';
 import { VisualizationsSetup } from '../../visualizations/public';
-import { ExpressionsStart } from '../../expressions/public';
+import { ExpressionsStart, ExpressionsPublicPlugin } from '../../expressions/public';
 import { NavigationPublicPluginStart } from '../../navigation/public';
 import { DataPublicPluginStart } from '../../data/public';
 import { TypeServiceSetup, TypeServiceStart } from './services/type_service';
@@ -28,6 +28,7 @@ export interface VisBuilderPluginSetupDependencies {
   embeddable: EmbeddableSetup;
   visualizations: VisualizationsSetup;
   data: DataPublicPluginSetup;
+  expressions: ReturnType<ExpressionsPublicPlugin['setup']>;
 }
 export interface VisBuilderPluginStartDependencies {
   embeddable: EmbeddableStart;

src/plugins/vis_builder/public/visualizations/common/expression_helpers.ts

@@ -7,9 +7,8 @@ import { cloneDeep } from 'lodash';
 import { OpenSearchaggsExpressionFunctionDefinition } from '../../../../data/public';
 import { ExpressionFunctionOpenSearchDashboards } from '../../../../expressions';
 import { buildExpressionFunction } from '../../../../expressions/public';
-import { VisualizationState } from '../../application/utils/state_management';
+import { VisualizationState, StyleState } from '../../application/utils/state_management';
 import { getSearchService, getIndexPatterns } from '../../plugin_services';
-import { StyleState } from '../../application/utils/state_management';
 
 export const getAggExpressionFunctions = async (
   visualization: VisualizationState,
@@ -49,3 +48,47 @@ export const getAggExpressionFunctions = async (
     expressionFns: [opensearchDashboards, opensearchaggs],
   };
 };
+
+export const getAggExpressionFunctionsWithContext = async (
+  visualization: VisualizationState,
+  style?: StyleState,
+  searchContext?: IExpressionLoaderParams['searchContext']
+) => {
+  const { activeVisualization, indexPattern: indexId = '' } = visualization;
+  const { aggConfigParams } = activeVisualization || {};
+
+  const indexPatternsService = getIndexPatterns();
+  const indexPattern = await indexPatternsService.get(indexId);
+  const aggConfigs = getSearchService().aggs.createAggConfigs(
+    indexPattern,
+    cloneDeep(aggConfigParams)
+  );
+
+  const opensearchDashboards = buildExpressionFunction<ExpressionFunctionOpenSearchDashboards>(
+    'opensearchDashboards',
+    {}
+  );
+
+  const opensearchDashboardsContext = buildExpressionFunction('opensearch_dashboards_context', {
+    timeRange: JSON.stringify(searchContext?.timeRange || {}),
+    filters: JSON.stringify(searchContext?.filters || []),
+    query: JSON.stringify(searchContext?.query || []),
+  });
+
+  const opensearchaggs = buildExpressionFunction<OpenSearchaggsExpressionFunctionDefinition>(
+    'opensearchaggs',
+    {
+      index: indexId,
+      metricsAtAllLevels: style?.showMetricsAtAllLevels || false,
+      partialRows: style?.showPartialRows || false,
+      aggConfigs: JSON.stringify(aggConfigs.aggs),
+      includeFormatHints: false,
+    }
+  );
+
+  return {
+    aggConfigs,
+    indexPattern,
+    expressionFns: [opensearchDashboards, opensearchDashboardsContext, opensearchaggs],
+  };
+};

src/plugins/vis_builder/public/visualizations/vega/build_spec_vega.ts

@@ -0,0 +1,126 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { buildEncoding } from './components/encoding';
+import { buildMark } from './components/mark';
+import { buildLegend } from './components/legend';
+import { VegaSpec, AxisFormats } from './utils/types';
+import { StyleState } from '../../application/utils/state_management';
+
+/**
+ * Builds a Vega specification based on the provided data, visual configuration, and style.
+ *
+ * @param {object} data - The data object containing series and axis information.
+ * @param {any} visConfig - The visual configuration settings.
+ * @param {StyleState} style - The style configuration for the visualization.
+ * @returns {VegaSpec} The complete Vega specification.
+ */
+export const buildVegaSpecViaVega = (data: any, visConfig: any, style: StyleState): VegaSpec => {
+  const { dimensions, addLegend, legendPosition } = visConfig;
+  const { type } = style;
+  const {
+    xAxisFormat,
+    xAxisLabel,
+    yAxisFormat,
+    yAxisLabel,
+    zAxisFormat,
+    series: transformedData,
+  } = data;
+
+  const formats: AxisFormats = {
+    xAxisFormat,
+    xAxisLabel,
+    yAxisFormat,
+    yAxisLabel,
+    zAxisFormat,
+  };
+
+  const spec: VegaSpec = {
+    $schema: 'https://vega.github.io/schema/vega/v5.json',
+    padding: 5,
+    data: [
+      {
+        name: 'source',
+        values: transformedData,
+      },
+      {
+        name: 'splits',
+        source: 'source',
+        transform: [
+          {
+            type: 'aggregate',
+            groupby: ['split'],
+          },
+        ],
+      },
+    ],
+    signals: [
+      { name: 'splitCount', update: 'length(data("splits"))' },
+      { name: 'chartWidth', update: 'width / splitCount - 10' },
+    ],
+    scales: [
+      {
+        name: 'splitScale',
+        type: 'band',
+        domain: { data: 'splits', field: 'split' },
+        range: 'width',
+        padding: 0.1,
+      },
+      {
+        name: 'color',
+        type: 'ordinal',
+        domain: { data: 'source', field: 'series' },
+        range: 'category',
+      },
+    ],
+    layout: {
+      columns: { signal: 'splitCount' },
+      padding: { row: 40, column: 20 },
+    },
+    marks: [
+      {
+        type: 'group',
+        from: { data: 'splits' },
+        encode: {
+          enter: {
+            width: { signal: 'chartWidth' },
+            height: { signal: 'height' },
+            stroke: { value: '#ccc' },
+            strokeWidth: { value: 1 },
+          },
+        },
+        signals: [{ name: 'width', update: 'chartWidth' }],
+        scales: buildEncoding(dimensions, formats, true),
+        axes: [
+          {
+            orient: 'bottom',
+            scale: 'xscale',
+            zindex: 1,
+            labelAngle: -90,
+            labelAlign: 'right',
+            labelBaseline: 'middle',
+          },
+          { orient: 'left', scale: 'yscale', zindex: 1 },
+        ],
+        title: {
+          text: { signal: 'parent.split' },
+          anchor: 'middle',
+          offset: 10,
+          limit: { signal: 'chartWidth' },
+          wrap: true,
+          align: 'center',
+        },
+        marks: buildMark(type, true),
+      },
+    ],
+  };
+
+  // Add legend if specified
+  if (addLegend) {
+    spec.legends = [buildLegend(legendPosition, true)];
+  }
+
+  return spec;
+};