[VisBuilder] Add Capability to generate dynamic vega (#7288)

* [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]>

* fix PR comments

* update file and functions names
* fix type errors
* fix area chart

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

* add unit tests

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

* enable embeddable for useVega

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

* remove buildVegaScales due to split it to smaller modules

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

* fix date for vega

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

* fix test

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

* Changeset file for PR #7288 created/updated

---------

Signed-off-by: Anan Zhuang <[email protected]>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>

changelogs/fragments/7288.yml

@@ -0,0 +1,2 @@
+feat:
+- [VisBuilder] Add Capability to generate dynamic vega ([#7288](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7288))

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/embeddable/vis_builder_embeddable.tsx

@@ -34,11 +34,13 @@ import {
   getIndexPatterns,
   getTypeService,
   getUIActions,
+  getUISettings,
 } from '../plugin_services';
 import { PersistedState, prepareJson } from '../../../visualizations/public';
 import { VisBuilderSavedVis } from '../saved_visualizations/transforms';
 import { handleVisEvent } from '../application/utils/handle_vis_event';
 import { VisBuilderEmbeddableFactoryDeps } from './vis_builder_embeddable_factory';
+import { VISBUILDER_ENABLE_VEGA_SETTING } from '../../common/constants';
 
 // Apparently this needs to match the saved object type for the clone and replace panel actions to work
 export const VISBUILDER_EMBEDDABLE = VISBUILDER_SAVED_OBJECT;
@@ -150,11 +152,16 @@ export class VisBuilderEmbeddable extends Embeddable<VisBuilderInput, VisBuilder
 
       if (!valid && errorMsg) throw new Error(errorMsg);
 
-      const exp = await toExpression(renderState, {
-        filters: this.filters,
-        query: this.query,
-        timeRange: this.timeRange,
-      });
+      const useVega = getUISettings().get(VISBUILDER_ENABLE_VEGA_SETTING);
+      const exp = await toExpression(
+        renderState,
+        {
+          filters: this.filters,
+          query: this.query,
+          timeRange: this.timeRange,
+        },
+        useVega
+      );
       return exp;
     } catch (error) {
       this.onContainerError(error as Error);

src/plugins/vis_builder/public/plugin.test.ts

@@ -9,6 +9,7 @@ import { dataPluginMock } from '../../data/public/mocks';
 import { embeddablePluginMock } from '../../embeddable/public/mocks';
 import { navigationPluginMock } from '../../navigation/public/mocks';
 import { visualizationsPluginMock } from '../../visualizations/public/mocks';
+import { expressionsPluginMock } from '../../expressions/public/mocks';
 import { PLUGIN_ID, PLUGIN_NAME } from '../common';
 import { VisBuilderPlugin } from './plugin';
 
@@ -29,6 +30,7 @@ describe('VisBuilderPlugin', () => {
         visualizations: visualizationsPluginMock.createSetupContract(),
         embeddable: embeddablePluginMock.createSetupContract(),
         data: dataPluginMock.createSetupContract(),
+        expressions: expressionsPluginMock.createSetupContract(), // Add this line
       };
 
       const setup = plugin.setup(coreSetup, setupDeps);
@@ -41,6 +43,7 @@ describe('VisBuilderPlugin', () => {
           aliasApp: PLUGIN_ID,
         })
       );
+      expect(setupDeps.expressions.registerFunction).toHaveBeenCalled(); // Add this expectation
     });
   });
 });

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';
@@ -18,6 +18,7 @@ import { IOsdUrlStateStorage } from '../../opensearch_dashboards_utils/public';
 import { DataPublicPluginSetup } from '../../data/public';
 import { UiActionsStart } from '../../ui_actions/public';
 import { Capabilities } from '../../../core/public';
+import { IUiSettingsClient } from '../../../core/public';
 
 export type VisBuilderSetup = TypeServiceSetup;
 export interface VisBuilderStart extends TypeServiceStart {
@@ -28,6 +29,7 @@ export interface VisBuilderPluginSetupDependencies {
   embeddable: EmbeddableSetup;
   visualizations: VisualizationsSetup;
   data: DataPublicPluginSetup;
+  expressions: ReturnType<ExpressionsPublicPlugin['setup']>;
 }
 export interface VisBuilderPluginStartDependencies {
   embeddable: EmbeddableStart;
@@ -37,6 +39,7 @@ export interface VisBuilderPluginStartDependencies {
   dashboard: DashboardStart;
   expressions: ExpressionsStart;
   uiActions: UiActionsStart;
+  uiSettings: IUiSettingsClient;
 }
 
 export interface VisBuilderServices extends CoreStart {

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

@@ -7,20 +7,21 @@ 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';
+import { IExpressionLoaderParams } from '../../../../expressions/public';
 
 export const getAggExpressionFunctions = async (
   visualization: VisualizationState,
-  style?: StyleState
+  style?: StyleState,
+  useVega: boolean = false,
+  searchContext?: IExpressionLoaderParams['searchContext']
 ) => {
   const { activeVisualization, indexPattern: indexId = '' } = visualization;
   const { aggConfigParams } = activeVisualization || {};
 
   const indexPatternsService = getIndexPatterns();
   const indexPattern = await indexPatternsService.get(indexId);
-  // aggConfigParams is the serealizeable aggConfigs that need to be reconstructed here using the agg servce
   const aggConfigs = getSearchService().aggs.createAggConfigs(
     indexPattern,
     cloneDeep(aggConfigParams)
@@ -31,7 +32,6 @@ export const getAggExpressionFunctions = async (
     {}
   );
 
-  // soon this becomes: const opensearchaggs = vis.data.aggs!.toExpressionAst();
   const opensearchaggs = buildExpressionFunction<OpenSearchaggsExpressionFunctionDefinition>(
     'opensearchaggs',
     {
@@ -43,9 +43,20 @@ export const getAggExpressionFunctions = async (
     }
   );
 
+  let expressionFns = [opensearchDashboards, opensearchaggs];
+
+  if (useVega === true && searchContext) {
+    const opensearchDashboardsContext = buildExpressionFunction('opensearch_dashboards_context', {
+      timeRange: JSON.stringify(searchContext.timeRange || {}),
+      filters: JSON.stringify(searchContext.filters || []),
+      query: JSON.stringify(searchContext.query || []),
+    });
+    expressionFns = [opensearchDashboards, opensearchDashboardsContext, opensearchaggs];
+  }
+
   return {
     aggConfigs,
     indexPattern,
-    expressionFns: [opensearchDashboards, opensearchaggs],
+    expressionFns,
   };
 };

src/plugins/vis_builder/public/visualizations/vega/components/axes.test.ts

@@ -0,0 +1,60 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { buildAxes } from './axes';
+
+describe('axes.ts', () => {
+  describe('buildAxes', () => {
+    it('should return correct axis configurations for date x-axis', () => {
+      const dimensions = {
+        x: { format: { id: 'date' } },
+        y: [{ label: 'Y Axis' }],
+      };
+      const formats = {
+        xAxisLabel: 'X Axis',
+        yAxisLabel: 'Custom Y Axis',
+      };
+
+      const result = buildAxes(dimensions, formats);
+
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({
+        orient: 'bottom',
+        scale: 'x',
+        labelAngle: -90,
+        labelAlign: 'right',
+        labelBaseline: 'middle',
+        title: 'X Axis',
+        format: '%Y-%m-%d %H:%M',
+      });
+      expect(result[1]).toEqual({
+        orient: 'left',
+        scale: 'y',
+        title: 'Custom Y Axis',
+      });
+    });
+
+    it('should not add format when x is not date', () => {
+      const dimensions = {
+        x: { format: { id: 'number' } },
+        y: [{ label: 'Y Axis' }],
+      };
+      const result = buildAxes(dimensions, 'X', 'Y');
+
+      expect(result[0]).not.toHaveProperty('format');
+    });
+
+    it('should use default labels when not provided', () => {
+      const dimensions = {
+        x: {},
+        y: [{ label: 'Default Y' }],
+      };
+      const result = buildAxes(dimensions, '', '');
+
+      expect(result[0].title).toBe('_all');
+      expect(result[1].title).toBe('Default Y');
+    });
+  });
+});

src/plugins/vis_builder/public/visualizations/vega/components/axes.ts

@@ -0,0 +1,51 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { AxisFormats } from '../utils/types';
+
+export interface AxisConfig {
+  orient?: string;
+  scale?: string;
+  labelAngle?: number;
+  labelAlign?: string;
+  labelBaseline?: string;
+  title: any;
+  format?: string; // property for date format
+}
+
+/**
+ * Builds the axes configuration for a chart.
+ *
+ * Note: This axis configuration is currently tailored for specific use cases.
+ * In the future, we plan to expand and generalize this function to accommodate
+ * a wider range of chart types and axis configurations.
+ * @param {any} dimensions - The dimensions of the data.
+ * @param {AxisFormats} formats - The formatting information for axes.
+ */
+
+export const buildAxes = (dimensions: any, formats: AxisFormats): AxisConfig[] => {
+  const { xAxisLabel, yAxisLabel } = formats;
+  const xAxis: AxisConfig = {
+    orient: 'bottom',
+    scale: 'x',
+    labelAngle: -90,
+    labelAlign: 'right',
+    labelBaseline: 'middle',
+    title: xAxisLabel || '_all',
+  };
+
+  // Add date format if x dimension is a date type
+  if (dimensions.x && dimensions.x.format && dimensions.x.format.id === 'date') {
+    xAxis.format = '%Y-%m-%d %H:%M';
+  }
+
+  const yAxis: AxisConfig = {
+    orient: 'left',
+    scale: 'y',
+    title: yAxisLabel ? yAxisLabel : dimensions.y && dimensions.y[0] ? dimensions.y[0].label : '',
+  };
+
+  return [xAxis, yAxis];
+};