chore(superset-ui-chart-controls): refactor pivot and rename operator (…

…#22963)
apache · Feb 2, 2023 · c53c3aa · c53c3aa
1 parent ed7b353
commit c53c3aa
Show file tree

Hide file tree

Showing 7 changed files with 100 additions and 18 deletions.
diff --git a/superset-frontend/packages/superset-ui-chart-controls/src/operators/pivotOperator.ts b/superset-frontend/packages/superset-ui-chart-controls/src/operators/pivotOperator.ts
@@ -31,13 +31,14 @@ export const pivotOperator: PostProcessingFactory<PostProcessingPivot> = (
 ) => {
   const metricLabels = ensureIsArray(queryObject.metrics).map(getMetricLabel);
   const xAxisLabel = getXAxisLabel(formData);
+  const columns = queryObject.series_columns || queryObject.columns;
 
   if (xAxisLabel && metricLabels.length) {
     return {
       operation: 'pivot',
       options: {
         index: [xAxisLabel],
-        columns: ensureIsArray(queryObject.columns).map(getColumnLabel),
+        columns: ensureIsArray(columns).map(getColumnLabel),
         // Create 'dummy' mean aggregates to assign cell values in pivot table
         // use the 'mean' aggregates to avoid drop NaN. PR: https://github.com/apache-superset/superset-ui/pull/1231
         aggregates: Object.fromEntries(

diff --git a/superset-frontend/packages/superset-ui-chart-controls/src/operators/renameOperator.ts b/superset-frontend/packages/superset-ui-chart-controls/src/operators/renameOperator.ts
@@ -32,14 +32,16 @@ export const renameOperator: PostProcessingFactory<PostProcessingRename> = (
   queryObject,
 ) => {
   const metrics = ensureIsArray(queryObject.metrics);
-  const columns = ensureIsArray(queryObject.columns);
+  const columns = ensureIsArray(
+    queryObject.series_columns || queryObject.columns,
+  );
   const { truncate_metric } = formData;
   const xAxisLabel = getXAxisLabel(formData);
   // remove or rename top level of column name(metric name) in the MultiIndex when
   // 1) only 1 metric
-  // 2) exist dimentsion
-  // 3) exist xAxis
-  // 4) exist time comparison, and comparison type is "actual values"
+  // 2) dimension exist
+  // 3) xAxis exist
+  // 4) time comparison exist, and comparison type is "actual values"
   // 5) truncate_metric in form_data and truncate_metric is true
   if (
     metrics.length === 1 &&

diff --git a/...et-frontend/packages/superset-ui-chart-controls/src/operators/timeComparePivotOperator.ts b/...et-frontend/packages/superset-ui-chart-controls/src/operators/timeComparePivotOperator.ts
@@ -31,6 +31,7 @@ export const timeComparePivotOperator: PostProcessingFactory<PostProcessingPivot
   (formData, queryObject) => {
     const metricOffsetMap = getMetricOffsetsMap(formData, queryObject);
     const xAxisLabel = getXAxisLabel(formData);
+    const columns = queryObject.series_columns || queryObject.columns;
 
     if (isTimeComparison(formData, queryObject) && xAxisLabel) {
       const aggregates = Object.fromEntries(
@@ -45,7 +46,7 @@ export const timeComparePivotOperator: PostProcessingFactory<PostProcessingPivot
         operation: 'pivot',
         options: {
           index: [xAxisLabel],
-          columns: ensureIsArray(queryObject.columns).map(getColumnLabel),
+          columns: ensureIsArray(columns).map(getColumnLabel),
           drop_missing_columns: !formData?.show_empty_columns,
           aggregates,
         },

diff --git a/superset-frontend/packages/superset-ui-chart-controls/test/operators/pivotOperator.test.ts b/superset-frontend/packages/superset-ui-chart-controls/test/operators/pivotOperator.test.ts
@@ -64,7 +64,7 @@ test('skip pivot', () => {
   ).toEqual(undefined);
 });
 
-test('pivot by __timestamp without groupby', () => {
+test('pivot by __timestamp without columns', () => {
   expect(
     pivotOperator(
       { ...formData, granularity_sqla: 'time_column' },
@@ -84,7 +84,7 @@ test('pivot by __timestamp without groupby', () => {
   });
 });
 
-test('pivot by __timestamp with groupby', () => {
+test('pivot by __timestamp with columns', () => {
   expect(
     pivotOperator(
       { ...formData, granularity_sqla: 'time_column' },
@@ -107,6 +107,29 @@ test('pivot by __timestamp with groupby', () => {
   });
 });
 
+test('pivot by __timestamp with series_columns', () => {
+  expect(
+    pivotOperator(
+      { ...formData, granularity_sqla: 'time_column' },
+      {
+        ...queryObject,
+        series_columns: ['foo', 'bar'],
+      },
+    ),
+  ).toEqual({
+    operation: 'pivot',
+    options: {
+      index: ['__timestamp'],
+      columns: ['foo', 'bar'],
+      aggregates: {
+        'count(*)': { operator: 'mean' },
+        'sum(val)': { operator: 'mean' },
+      },
+      drop_missing_columns: false,
+    },
+  });
+});
+
 test('pivot by x_axis with groupby', () => {
   expect(
     pivotOperator(
@@ -116,7 +139,7 @@ test('pivot by x_axis with groupby', () => {
       },
       {
         ...queryObject,
-        columns: ['foo', 'bar'],
+        series_columns: ['foo', 'bar'],
       },
     ),
   ).toEqual({
@@ -146,7 +169,7 @@ test('pivot by adhoc x_axis', () => {
       },
       {
         ...queryObject,
-        columns: ['foo', 'bar'],
+        series_columns: ['foo', 'bar'],
       },
     ),
   ).toEqual({

diff --git a/superset-frontend/packages/superset-ui-chart-controls/test/operators/renameOperator.test.ts b/superset-frontend/packages/superset-ui-chart-controls/test/operators/renameOperator.test.ts
@@ -49,7 +49,7 @@ test('should skip renameOperator if exists multiple metrics', () => {
   ).toEqual(undefined);
 });
 
-test('should skip renameOperator if does not exist series', () => {
+test('should skip renameOperator if series does not exist', () => {
   expect(
     renameOperator(formData, {
       ...queryObject,
@@ -105,7 +105,7 @@ test('should add renameOperator', () => {
   });
 });
 
-test('should add renameOperator if does not exist x_axis', () => {
+test('should add renameOperator if x_axis does not exist', () => {
   expect(
     renameOperator(
       {
@@ -120,6 +120,25 @@ test('should add renameOperator if does not exist x_axis', () => {
   });
 });
 
+test('should add renameOperator if based on series_columns', () => {
+  expect(
+    renameOperator(
+      {
+        ...formData,
+        ...{ x_axis: null, granularity_sqla: 'time column' },
+      },
+      {
+        ...queryObject,
+        columns: [],
+        series_columns: ['gender', 'dttm'],
+      },
+    ),
+  ).toEqual({
+    operation: 'rename',
+    options: { columns: { 'count(*)': null }, inplace: true, level: 0 },
+  });
+});
+
 test('should add renameOperator if exist "actual value" time comparison', () => {
   expect(
     renameOperator(

diff --git a/...ntend/packages/superset-ui-chart-controls/test/operators/timeComparePivotOperator.test.ts b/...ntend/packages/superset-ui-chart-controls/test/operators/timeComparePivotOperator.test.ts
@@ -136,6 +136,45 @@ test('should pivot on x-axis', () => {
   });
 });
 
+test('should pivot on x-axis with series_columns', () => {
+  expect(
+    timeComparePivotOperator(
+      {
+        ...formData,
+        comparison_type: 'values',
+        time_compare: ['1 year ago', '1 year later'],
+        x_axis: 'ds',
+      },
+      {
+        ...queryObject,
+        columns: ['ds', 'foo', 'bar'],
+        series_columns: ['foo', 'bar'],
+      },
+    ),
+  ).toEqual({
+    operation: 'pivot',
+    options: {
+      aggregates: {
+        'count(*)': { operator: 'mean' },
+        'count(*)__1 year ago': { operator: 'mean' },
+        'count(*)__1 year later': { operator: 'mean' },
+        'sum(val)': {
+          operator: 'mean',
+        },
+        'sum(val)__1 year ago': {
+          operator: 'mean',
+        },
+        'sum(val)__1 year later': {
+          operator: 'mean',
+        },
+      },
+      drop_missing_columns: false,
+      columns: ['foo', 'bar'],
+      index: ['ds'],
+    },
+  });
+});
+
 test('should pivot on adhoc x-axis', () => {
   expect(
     timeComparePivotOperator(

diff --git a/superset-frontend/plugins/plugin-chart-echarts/src/MixedTimeseries/buildQuery.ts b/superset-frontend/plugins/plugin-chart-echarts/src/MixedTimeseries/buildQuery.ts
@@ -67,8 +67,8 @@ export default function buildQuery(formData: QueryFormData) {
         fd,
         queryObject,
       )
-        ? timeComparePivotOperator(fd, { ...queryObject, columns: fd.groupby })
-        : pivotOperator(fd, { ...queryObject, columns: fd.groupby });
+        ? timeComparePivotOperator(fd, queryObject)
+        : pivotOperator(fd, queryObject);
 
       const tmpQueryObject = {
         ...queryObject,
@@ -78,10 +78,7 @@ export default function buildQuery(formData: QueryFormData) {
           rollingWindowOperator(fd, queryObject),
           timeCompareOperator(fd, queryObject),
           resampleOperator(fd, queryObject),
-          renameOperator(fd, {
-            ...queryObject,
-            columns: fd.groupby,
-          }),
+          renameOperator(fd, queryObject),
           flattenOperator(fd, queryObject),
         ],
       } as QueryObject;