feat(plugin-chart-echarts): add support for custom forecasts (#1389)

apache · Nov 26, 2021 · f1613fd · f1613fd
1 parent 50f9d84
commit f1613fd
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 8 deletions.
diff --git a/...ary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/Timeseries/transformProps.ts b/...ary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/Timeseries/transformProps.ts
@@ -46,6 +46,7 @@ import {
 import { extractAnnotationLabels } from '../utils/annotation';
 import {
   extractForecastSeriesContext,
+  extractForecastSeriesContexts,
   extractProphetValuesFromTooltipParams,
   formatProphetTooltipSeries,
   rebaseTimeseriesDatum,
@@ -113,6 +114,9 @@ export default function transformProps(
   const rawSeries = extractTimeseriesSeries(rebasedData, {
     fillNeighborValue: stack && !forecastEnabled ? 0 : undefined,
   });
+  const seriesContexts = extractForecastSeriesContexts(
+    Object.values(rawSeries).map(series => series.name as string),
+  );
   const series: SeriesOption[] = [];
   const formatter = getNumberFormatter(contributionMode ? ',.0%' : yAxisFormat);
 
@@ -145,7 +149,7 @@ export default function transformProps(
     const transformedSeries = transformSeries(entry, colorScale, {
       area,
       filterState,
-      forecastEnabled,
+      seriesContexts,
       markerEnabled,
       markerSize,
       areaOpacity: opacity,

diff --git a/...orary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/Timeseries/transformers.ts b/...orary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/Timeseries/transformers.ts
@@ -67,7 +67,7 @@ export function transformSeries(
   opts: {
     area?: boolean;
     filterState?: FilterState;
-    forecastEnabled?: boolean;
+    seriesContexts?: { [key: string]: ForecastSeriesEnum[] };
     markerEnabled?: boolean;
     markerSize?: number;
     areaOpacity?: number;
@@ -86,7 +86,7 @@ export function transformSeries(
   const {
     area,
     filterState,
-    forecastEnabled,
+    seriesContexts = {},
     markerEnabled,
     markerSize,
     areaOpacity = 1,
@@ -100,6 +100,11 @@ export function transformSeries(
     showValueIndexes = [],
     richTooltip,
   } = opts;
+  const contexts = seriesContexts[name || ''] || [];
+  const hasForecast =
+    contexts.includes(ForecastSeriesEnum.ForecastTrend) ||
+    contexts.includes(ForecastSeriesEnum.ForecastLower) ||
+    contexts.includes(ForecastSeriesEnum.ForecastUpper);
 
   const forecastSeries = extractForecastSeriesContext(name || '');
   const isConfidenceBand =
@@ -125,7 +130,7 @@ export function transformSeries(
     stackId = forecastSeries.type;
   }
   let plotType;
-  if (!isConfidenceBand && (seriesType === 'scatter' || (forecastEnabled && isObservation))) {
+  if (!isConfidenceBand && (seriesType === 'scatter' || (hasForecast && isObservation))) {
     plotType = 'scatter';
   } else if (isConfidenceBand) {
     plotType = 'line';
@@ -141,7 +146,7 @@ export function transformSeries(
   if (!isConfidenceBand) {
     if (plotType === 'scatter') {
       showSymbol = true;
-    } else if (forecastEnabled && isObservation) {
+    } else if (hasForecast && isObservation) {
       showSymbol = true;
     } else if (plotType === 'line' && showValue) {
       showSymbol = true;

diff --git a/...ntend/temporary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/utils/prophet.ts b/...ntend/temporary_superset_ui/superset-ui/plugins/plugin-chart-echarts/src/utils/prophet.ts
@@ -35,6 +35,16 @@ export const extractForecastSeriesContext = (seriesName: OptionName): ForecastSe
   };
 };
 
+export const extractForecastSeriesContexts = (
+  seriesNames: string[],
+): { [key: string]: ForecastSeriesEnum[] } =>
+  seriesNames.reduce((agg, name) => {
+    const context = extractForecastSeriesContext(name);
+    const currentContexts = agg[context.name] || [];
+    currentContexts.push(context.type);
+    return { ...agg, [context.name]: currentContexts };
+  }, {} as { [key: string]: ForecastSeriesEnum[] });
+
 export const extractProphetValuesFromTooltipParams = (
   params: (CallbackDataParams & { seriesId: string })[],
 ): Record<string, ProphetValue> => {
@@ -83,10 +93,12 @@ export const formatProphetTooltipSeries = ({
   if (forecastTrend) {
     if (isObservation) row += ', ';
     row += `ŷ = ${formatter(forecastTrend)}`;
-    if (forecastLower && forecastUpper)
-      // the lower bound needs to be added to the upper bound
-      row += ` (${formatter(forecastLower)}, ${formatter(forecastLower + forecastUpper)})`;
   }
+  if (forecastLower && forecastUpper)
+    // the lower bound needs to be added to the upper bound
+    row = `${row.trim()} (${formatter(forecastLower)}, ${formatter(
+      forecastLower + forecastUpper,
+    )})`;
   return `${row.trim()}`;
 };
 

diff --git a/...temporary_superset_ui/superset-ui/plugins/plugin-chart-echarts/test/utils/prophet.test.ts b/...temporary_superset_ui/superset-ui/plugins/plugin-chart-echarts/test/utils/prophet.test.ts
@@ -171,5 +171,24 @@ describe('formatProphetTooltipSeries', () => {
         formatter,
       }),
     ).toEqual('<img>qwerty: ŷ = 20 (5, 12)');
+    expect(
+      formatProphetTooltipSeries({
+        seriesName: 'qwerty',
+        marker: '<img>',
+        observation: 10.1,
+        forecastLower: 6,
+        forecastUpper: 7,
+        formatter,
+      }),
+    ).toEqual('<img>qwerty: 10 (6, 13)');
+    expect(
+      formatProphetTooltipSeries({
+        seriesName: 'qwerty',
+        marker: '<img>',
+        forecastLower: 7,
+        forecastUpper: 8,
+        formatter,
+      }),
+    ).toEqual('<img>qwerty: (7, 15)');
   });
 });