fix(plugin-chart-echarts): support forced categorical x-axis (#26404)

superset-frontend/packages/superset-ui-chart-controls/src/constants.ts

@@ -17,12 +17,11 @@
  * under the License.
  */
 import {
-  t,
-  QueryMode,
   DTTM_ALIAS,
   GenericDataType,
   QueryColumn,
-  DatasourceType,
+  QueryMode,
+  t,
 } from '@superset-ui/core';
 import { ColumnMeta, SortSeriesData, SortSeriesType } from './types';
 
@@ -43,6 +42,7 @@ export const COLUMN_NAME_ALIASES: Record<string, string> = {
 export const DATASET_TIME_COLUMN_OPTION: ColumnMeta = {
   verbose_name: COLUMN_NAME_ALIASES[DTTM_ALIAS],
   column_name: DTTM_ALIAS,
+  type: 'TIMESTAMP',
   type_generic: GenericDataType.TEMPORAL,
   description: t(
     'A reference to the [Time] configuration, taking granularity into account',
@@ -51,8 +51,9 @@ export const DATASET_TIME_COLUMN_OPTION: ColumnMeta = {
 
 export const QUERY_TIME_COLUMN_OPTION: QueryColumn = {
   column_name: DTTM_ALIAS,
-  type: DatasourceType.Query,
+  is_dttm: true,
-  is_dttm: false,
+  type: 'TIMESTAMP',
+  type_generic: GenericDataType.TEMPORAL,
 };
 
 export const QueryModeLabel = {

superset-frontend/packages/superset-ui-chart-controls/src/fixtures.ts

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { DatasourceType } from '@superset-ui/core';
+import { DatasourceType, GenericDataType } from '@superset-ui/core';
 import { Dataset } from './types';
 
 export const TestDataset: Dataset = {
@@ -37,7 +37,7 @@ export const TestDataset: Dataset = {
       is_dttm: false,
       python_date_format: null,
       type: 'BIGINT',
-      type_generic: 0,
+      type_generic: GenericDataType.NUMERIC,
       verbose_name: null,
       warning_markdown: null,
     },
@@ -55,7 +55,7 @@ export const TestDataset: Dataset = {
       is_dttm: false,
       python_date_format: null,
       type: 'VARCHAR(16)',
-      type_generic: 1,
+      type_generic: GenericDataType.STRING,
       verbose_name: '',
       warning_markdown: null,
     },
@@ -73,7 +73,7 @@ export const TestDataset: Dataset = {
       is_dttm: false,
       python_date_format: null,
       type: 'VARCHAR(10)',
-      type_generic: 1,
+      type_generic: GenericDataType.STRING,
       verbose_name: null,
       warning_markdown: null,
     },
@@ -91,7 +91,7 @@ export const TestDataset: Dataset = {
       is_dttm: true,
       python_date_format: null,
       type: 'TIMESTAMP WITHOUT TIME ZONE',
-      type_generic: 2,
+      type_generic: GenericDataType.TEMPORAL,
       verbose_name: null,
       warning_markdown: null,
     },
@@ -109,7 +109,7 @@ export const TestDataset: Dataset = {
       is_dttm: false,
       python_date_format: null,
       type: 'VARCHAR(255)',
-      type_generic: 1,
+      type_generic: GenericDataType.STRING,
       verbose_name: null,
       warning_markdown: null,
     },

superset-frontend/packages/superset-ui-chart-controls/src/sections/echartsTimeSeriesQuery.tsx

@@ -20,6 +20,7 @@ import { hasGenericChartAxes, t } from '@superset-ui/core';
 import { ControlPanelSectionConfig, ControlSetRow } from '../types';
 import {
   contributionModeControl,
+  xAxisForceCategoricalControl,
   xAxisSortAscControl,
   xAxisSortControl,
   xAxisSortSeriesAscendingControl,
@@ -55,6 +56,7 @@ export const echartsTimeSeriesQueryWithXAxisSort: ControlPanelSectionConfig = {
   controlSetRows: [
     [hasGenericChartAxes ? 'x_axis' : null],
     [hasGenericChartAxes ? 'time_grain_sqla' : null],
+    [hasGenericChartAxes ? xAxisForceCategoricalControl : null],
     [hasGenericChartAxes ? xAxisSortControl : null],
     [hasGenericChartAxes ? xAxisSortAscControl : null],
     [hasGenericChartAxes ? xAxisSortSeriesControl : null],

superset-frontend/packages/superset-ui-chart-controls/src/shared-controls/customControls.tsx

@@ -20,9 +20,9 @@
 import {
   ContributionType,
   ensureIsArray,
+  GenericDataType,
   getColumnLabel,
   getMetricLabel,
-  isDefined,
   QueryFormColumn,
   QueryFormMetric,
   t,
@@ -38,6 +38,7 @@ import {
   DEFAULT_XAXIS_SORT_SERIES_DATA,
   SORT_SERIES_CHOICES,
 } from '../constants';
+import { checkColumnType } from '../utils/checkColumnType';
 
 export const contributionModeControl = {
   name: 'contributionMode',
@@ -54,18 +55,29 @@ export const contributionModeControl = {
   },
 };
 
-function isTemporal(controls: ControlStateMapping): boolean {
+function isForcedCategorical(controls: ControlStateMapping): boolean {
-  return !(
+  return (
-    isDefined(controls?.x_axis?.value) &&
+    checkColumnType(
-    !isTemporalColumn(
       getColumnLabel(controls?.x_axis?.value as QueryFormColumn),
       controls?.datasource?.datasource,
+      [GenericDataType.NUMERIC],
+    ) && !!controls?.xAxisForceCategorical?.value
+  );
+}
+
+function isSortable(controls: ControlStateMapping): boolean {
+  return (
+    isForcedCategorical(controls) ||
+    checkColumnType(
+      getColumnLabel(controls?.x_axis?.value as QueryFormColumn),
+      controls?.datasource?.datasource,
+      [GenericDataType.STRING, GenericDataType.BOOLEAN],
     )
   );
 }
 
 const xAxisSortVisibility = ({ controls }: { controls: ControlStateMapping }) =>
-  !isTemporal(controls) &&
+  isSortable(controls) &&
   ensureIsArray(controls?.groupby?.value).length === 0 &&
   ensureIsArray(controls?.metrics?.value).length === 1;
 
@@ -74,7 +86,7 @@ const xAxisMultiSortVisibility = ({
 }: {
   controls: ControlStateMapping;
 }) =>
-  !isTemporal(controls) &&
+  isSortable(controls) &&
   (!!ensureIsArray(controls?.groupby?.value).length ||
     ensureIsArray(controls?.metrics?.value).length > 1);
 
@@ -141,7 +153,29 @@ export const xAxisSortAscControl = {
         : t('X-Axis Sort Ascending'),
     default: true,
     description: t('Whether to sort ascending or descending on the base Axis.'),
-    visibility: xAxisSortVisibility,
+    visibility: ({ controls }: { controls: ControlStateMapping }) =>
+      controls?.x_axis_sort?.value !== undefined &&
+      xAxisSortVisibility({ controls }),
+  },
+};
+
+export const xAxisForceCategoricalControl = {
+  name: 'xAxisForceCategorical',
+  config: {
+    type: 'CheckboxControl',
+    label: () => t('Force categorical'),
+    default: false,
+    description: t('Treat values as categorical.'),
+    initialValue: (control: ControlState, state: ControlPanelState | null) =>
+      state?.form_data?.x_axis_sort !== undefined || control.value,
+    renderTrigger: true,
+    visibility: ({ controls }: { controls: ControlStateMapping }) =>
+      checkColumnType(
+        getColumnLabel(controls?.x_axis?.value as QueryFormColumn),
+        controls?.datasource?.datasource,
+        [GenericDataType.NUMERIC],
+      ),
+    shouldMapStateToProps: () => true,
   },
 };
 
@@ -173,6 +207,8 @@ export const xAxisSortSeriesAscendingControl = {
     default: DEFAULT_XAXIS_SORT_SERIES_DATA.sort_series_ascending,
     description: t('Whether to sort ascending or descending on the base Axis.'),
     renderTrigger: true,
-    visibility: xAxisMultiSortVisibility,
+    visibility: ({ controls }: { controls: ControlStateMapping }) =>
+      controls?.x_axis_sort_series?.value !== undefined &&
+      xAxisMultiSortVisibility({ controls }),
   },
 };

superset-frontend/packages/superset-ui-chart-controls/src/types.ts

@@ -479,13 +479,15 @@ export function isControlPanelSectionConfig(
 export function isDataset(
   datasource: Dataset | QueryResponse | null | undefined,
 ): datasource is Dataset {
-  return !!datasource && 'columns' in datasource;
+  return (
+    !!datasource && 'columns' in datasource && !('sqlEditorId' in datasource)
+  );
 }
 
 export function isQueryResponse(
   datasource: Dataset | QueryResponse | null | undefined,
 ): datasource is QueryResponse {
-  return !!datasource && 'results' in datasource && 'sql' in datasource;
+  return !!datasource && 'results' in datasource && 'sqlEditorId' in datasource;
 }
 
 export enum SortSeriesType {

superset-frontend/packages/superset-ui-chart-controls/src/utils/checkColumnType.ts

@@ -0,0 +1,49 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { ensureIsArray, GenericDataType, ValueOf } from '@superset-ui/core';
+import {
+  ControlPanelState,
+  isDataset,
+  isQueryResponse,
+} from '@superset-ui/chart-controls';
+
+export function checkColumnType(
+  columnName: string,
+  datasource: ValueOf<Pick<ControlPanelState, 'datasource'>>,
+  columnTypes: GenericDataType[],
+): boolean {
+  if (isDataset(datasource)) {
+    return ensureIsArray(datasource.columns).some(
+      c =>
+        c.type_generic !== undefined &&
+        columnTypes.includes(c.type_generic) &&
+        columnName === c.column_name,
+    );
+  }
+  if (isQueryResponse(datasource)) {
+    return ensureIsArray(datasource.columns)
+      .filter(
+        c =>
+          c.type_generic !== undefined && columnTypes.includes(c.type_generic),
+      )
+      .map(c => c.column_name)
+      .some(c => columnName === c);
+  }
+  return false;
+}

superset-frontend/packages/superset-ui-chart-controls/src/utils/columnChoices.ts

@@ -17,7 +17,7 @@
  * under the License.
  */
 import { QueryResponse } from '@superset-ui/core';
-import { Dataset, isColumnMeta, isDataset } from '../types';
+import { Dataset, isDataset, isQueryResponse } from '../types';
 
 /**
  * Convert Datasource columns to column choices
@@ -25,11 +25,13 @@ import { Dataset, isColumnMeta, isDataset } from '../types';
 export default function columnChoices(
   datasource?: Dataset | QueryResponse | null,
 ): [string, string][] {
-  if (isDataset(datasource) && isColumnMeta(datasource.columns[0])) {
+  if (isDataset(datasource) || isQueryResponse(datasource)) {
     return datasource.columns
       .map((col): [string, string] => [
         col.column_name,
-        col.verbose_name || col.column_name,
+        'verbose_name' in col
+          ? col.verbose_name || col.column_name
+          : col.column_name,
       ])
       .sort((opt1, opt2) =>
         opt1[1].toLowerCase() > opt2[1].toLowerCase() ? 1 : -1,

superset-frontend/packages/superset-ui-chart-controls/src/utils/index.ts

@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+export * from './checkColumnType';
 export * from './selectOptions';
 export * from './D3Formatting';
 export * from './expandControlConfig';

superset-frontend/packages/superset-ui-chart-controls/test/utils/checkColumnType.test.ts

@@ -0,0 +1,48 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import { GenericDataType, testQueryResponse } from '@superset-ui/core';
+import { checkColumnType, TestDataset } from '../../src';
+
+test('checkColumnType columns from a Dataset', () => {
+  expect(
+    checkColumnType('num', TestDataset, [GenericDataType.NUMERIC]),
+  ).toEqual(true);
+  expect(checkColumnType('num', TestDataset, [GenericDataType.STRING])).toEqual(
+    false,
+  );
+  expect(
+    checkColumnType('gender', TestDataset, [GenericDataType.STRING]),
+  ).toEqual(true);
+  expect(
+    checkColumnType('gender', TestDataset, [GenericDataType.NUMERIC]),
+  ).toEqual(false);
+});
+
+test('checkColumnType from a QueryResponse', () => {
+  expect(
+    checkColumnType('Column 1', testQueryResponse, [GenericDataType.STRING]),
+  ).toEqual(true);
+  expect(
+    checkColumnType('Column 1', testQueryResponse, [GenericDataType.NUMERIC]),
+  ).toEqual(false);
+});
+
+test('checkColumnType from null', () => {
+  expect(checkColumnType('col', null, [])).toEqual(false);
+});

superset-frontend/packages/superset-ui-chart-controls/test/utils/columnChoices.test.tsx

@@ -16,7 +16,11 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { DatasourceType, testQueryResponse } from '@superset-ui/core';
+import {
+  DatasourceType,
+  GenericDataType,
+  testQueryResponse,
+} from '@superset-ui/core';
 import { columnChoices } from '../../src';
 
 describe('columnChoices()', () => {
@@ -31,14 +35,20 @@ describe('columnChoices()', () => {
         columns: [
           {
             column_name: 'fiz',
+            type: 'INT',
+            type_generic: GenericDataType.NUMERIC,
           },
           {
             column_name: 'about',
             verbose_name: 'right',
+            type: 'VARCHAR',
+            type_generic: GenericDataType.STRING,
           },
           {
             column_name: 'foo',
-            verbose_name: 'bar',
+            verbose_name: undefined,
+            type: 'TIMESTAMP',
+            type_generic: GenericDataType.TEMPORAL,
           },
         ],
         verbose_map: {},
@@ -48,8 +58,8 @@ describe('columnChoices()', () => {
         description: 'this is my datasource',
       }),
     ).toEqual([
-      ['foo', 'bar'],
       ['fiz', 'fiz'],
+      ['foo', 'foo'],
       ['about', 'right'],
     ]);
   });

superset-frontend/packages/superset-ui-chart-controls/test/utils/getTemporalColumns.test.ts

@@ -16,7 +16,11 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { testQueryResponse, testQueryResults } from '@superset-ui/core';
+import {
+  GenericDataType,
+  testQueryResponse,
+  testQueryResults,
+} from '@superset-ui/core';
 import {
   Dataset,
   getTemporalColumns,
@@ -55,8 +59,9 @@ test('get temporal columns from a QueryResponse', () => {
     temporalColumns: [
       {
         column_name: 'Column 2',
-        type: 'TIMESTAMP',
         is_dttm: true,
+        type: 'TIMESTAMP',
+        type_generic: GenericDataType.TEMPORAL,
       },
     ],
     defaultTemporalColumn: 'Column 2',

superset-frontend/packages/superset-ui-core/src/query/types/Query.ts

@@ -31,6 +31,7 @@ import { Maybe } from '../../types';
 import { PostProcessingRule } from './PostProcessing';
 import { JsonObject } from '../../connection';
 import { TimeGranularity } from '../../time-format';
+import { GenericDataType } from './QueryResponse';
 
 export type BaseQueryObjectFilterClause = {
   col: QueryFormColumn;
@@ -250,6 +251,7 @@ export type QueryColumn = {
   name?: string;
   column_name: string;
   type: string | null;
+  type_generic: GenericDataType;
   is_dttm: boolean;
 };
 
@@ -383,16 +385,19 @@ export const testQuery: Query = {
       column_name: 'Column 1',
       type: 'STRING',
       is_dttm: false,
+      type_generic: GenericDataType.STRING,
     },
     {
       column_name: 'Column 3',
       type: 'STRING',
       is_dttm: false,
+      type_generic: GenericDataType.STRING,
     },
     {
       column_name: 'Column 2',
       type: 'TIMESTAMP',
       is_dttm: true,
+      type_generic: GenericDataType.TEMPORAL,
     },
   ],
 };
@@ -404,16 +409,19 @@ export const testQueryResults = {
       {
         column_name: 'Column 1',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
         is_dttm: false,
       },
       {
         column_name: 'Column 3',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
         is_dttm: false,
       },
       {
         column_name: 'Column 2',
         type: 'TIMESTAMP',
+        type_generic: GenericDataType.TEMPORAL,
         is_dttm: true,
       },
     ],
@@ -425,16 +433,19 @@ export const testQueryResults = {
       {
         column_name: 'Column 1',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
         is_dttm: false,
       },
       {
         column_name: 'Column 3',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
         is_dttm: false,
       },
       {
         column_name: 'Column 2',
         type: 'TIMESTAMP',
+        type_generic: GenericDataType.TEMPORAL,
         is_dttm: true,
       },
     ],

superset-frontend/plugins/plugin-chart-echarts/src/MixedTimeseries/transformProps.ts

@@ -189,6 +189,7 @@ export default function transformProps(
     groupby,
     groupbyB,
     xAxis: xAxisOrig,
+    xAxisForceCategorical,
     xAxisTitle,
     yAxisTitle,
     xAxisTitleMargin,
@@ -227,7 +228,7 @@ export default function transformProps(
 
   const dataTypes = getColtypesMapping(queriesData[0]);
   const xAxisDataType = dataTypes?.[xAxisLabel] ?? dataTypes?.[xAxisOrig];
-  const xAxisType = getAxisType(stack, xAxisDataType);
+  const xAxisType = getAxisType(stack, xAxisForceCategorical, xAxisDataType);
   const series: SeriesOption[] = [];
   const formatter = contributionMode
     ? getNumberFormatter(',.0%')

superset-frontend/plugins/plugin-chart-echarts/src/MixedTimeseries/types.ts

@@ -106,6 +106,7 @@ export const DEFAULT_FORM_DATA: EchartsMixedTimeseriesFormData = {
   yAxisTitleSecondary: DEFAULT_TITLE_FORM_DATA.yAxisTitle,
   tooltipTimeFormat: TIMESERIES_DEFAULTS.tooltipTimeFormat,
   xAxisBounds: TIMESERIES_DEFAULTS.xAxisBounds,
+  xAxisForceCategorical: TIMESERIES_DEFAULTS.xAxisForceCategorical,
   xAxisTimeFormat: TIMESERIES_DEFAULTS.xAxisTimeFormat,
   area: TIMESERIES_DEFAULTS.area,
   areaB: TIMESERIES_DEFAULTS.area,

superset-frontend/plugins/plugin-chart-echarts/src/Timeseries/constants.ts

@@ -63,6 +63,7 @@ export const DEFAULT_FORM_DATA: EchartsTimeseriesFormData = {
   yAxisBounds: [null, null],
   zoomable: false,
   richTooltip: true,
+  xAxisForceCategorical: false,
   xAxisLabelRotation: defaultXAxis.xAxisLabelRotation,
   groupby: [],
   showValue: false,

superset-frontend/plugins/plugin-chart-echarts/src/Timeseries/transformProps.ts

@@ -167,6 +167,7 @@ export default function transformProps(
     truncateYAxis,
     xAxis: xAxisOrig,
     xAxisBounds,
+    xAxisForceCategorical,
     xAxisLabelRotation,
     xAxisSortSeries,
     xAxisSortSeriesAscending,
@@ -248,7 +249,7 @@ export default function transformProps(
   const isAreaExpand = stack === StackControlsValue.Expand;
   const xAxisDataType = dataTypes?.[xAxisLabel] ?? dataTypes?.[xAxisOrig];
 
-  const xAxisType = getAxisType(stack, xAxisDataType);
+  const xAxisType = getAxisType(stack, xAxisForceCategorical, xAxisDataType);
   const series: SeriesOption[] = [];
 
   const forcePercentFormatter = Boolean(contributionMode || isAreaExpand);

superset-frontend/plugins/plugin-chart-echarts/src/Timeseries/types.ts

@@ -79,6 +79,7 @@ export type EchartsTimeseriesFormData = QueryFormData & {
   truncateXAxis: boolean;
   truncateYAxis: boolean;
   yAxisFormat?: string;
+  xAxisForceCategorical?: boolean;
   xAxisTimeFormat?: string;
   timeGrainSqla?: TimeGranularity;
   xAxisBounds: [number | undefined | null, number | undefined | null];

superset-frontend/plugins/plugin-chart-echarts/src/controls.tsx

@@ -309,3 +309,14 @@ export const minorTicks: ControlSetItem = {
     description: t('Show minor ticks on axes.'),
   },
 };
+
+export const forceCategorical: ControlSetItem = {
+  name: 'forceCategorical',
+  config: {
+    type: 'CheckboxControl',
+    label: t('Force categorical'),
+    default: false,
+    renderTrigger: true,
+    description: t('Make the x-axis categorical'),
+  },
+};

superset-frontend/plugins/plugin-chart-echarts/src/utils/series.ts

@@ -229,8 +229,10 @@ export function sortRows(
     }
 
     const value =
-      xAxisSortSeries === SortSeriesType.Name && typeof sortKey === 'string'
+      xAxisSortSeries === SortSeriesType.Name
-        ? sortKey.toLowerCase()
+        ? typeof sortKey === 'string'
+          ? sortKey.toLowerCase()
+          : sortKey
         : aggregate;
 
     return {
@@ -515,8 +517,12 @@ export function sanitizeHtml(text: string): string {
 
 export function getAxisType(
   stack: StackType,
+  forceCategorical?: boolean,
   dataType?: GenericDataType,
 ): AxisType {
+  if (forceCategorical) {
+    return AxisType.category;
+  }
   if (dataType === GenericDataType.TEMPORAL) {
     return AxisType.time;
   }

superset-frontend/plugins/plugin-chart-echarts/test/utils/series.test.ts

@@ -878,14 +878,28 @@ test('calculateLowerLogTick', () => {
   expect(calculateLowerLogTick(0.005)).toEqual(0.001);
 });
 
-test('getAxisType', () => {
+test('getAxisType without forced categorical', () => {
-  expect(getAxisType(false, GenericDataType.TEMPORAL)).toEqual(AxisType.time);
+  expect(getAxisType(false, false, GenericDataType.TEMPORAL)).toEqual(
-  expect(getAxisType(false, GenericDataType.NUMERIC)).toEqual(AxisType.value);
+    AxisType.time,
-  expect(getAxisType(true, GenericDataType.NUMERIC)).toEqual(AxisType.category);
+  );
-  expect(getAxisType(false, GenericDataType.BOOLEAN)).toEqual(
+  expect(getAxisType(false, false, GenericDataType.NUMERIC)).toEqual(
+    AxisType.value,
+  );
+  expect(getAxisType(true, false, GenericDataType.NUMERIC)).toEqual(
+    AxisType.category,
+  );
+  expect(getAxisType(false, false, GenericDataType.BOOLEAN)).toEqual(
+    AxisType.category,
+  );
+  expect(getAxisType(false, false, GenericDataType.STRING)).toEqual(
+    AxisType.category,
+  );
+});
+
+test('getAxisType with forced categorical', () => {
+  expect(getAxisType(false, true, GenericDataType.NUMERIC)).toEqual(
     AxisType.category,
   );
-  expect(getAxisType(false, GenericDataType.STRING)).toEqual(AxisType.category);
 });
 
 test('getMinAndMaxFromBounds returns empty object when not truncating', () => {

superset-frontend/src/SqlLab/fixtures.ts

@@ -22,6 +22,7 @@ import { ColumnKeyTypeType } from 'src/SqlLab/components/ColumnElement';
 import {
   DatasourceType,
   denormalizeTimestamp,
+  GenericDataType,
   QueryResponse,
   QueryState,
 } from '@superset-ui/core';
@@ -581,11 +582,13 @@ const baseQuery: QueryResponse = {
         is_dttm: true,
         column_name: 'ds',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
       },
       {
         is_dttm: false,
         column_name: 'gender',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
       },
     ],
     selected_columns: [
@@ -593,11 +596,13 @@ const baseQuery: QueryResponse = {
         is_dttm: true,
         column_name: 'ds',
         type: 'STRING',
+        type_generic: GenericDataType.TEMPORAL,
       },
       {
         is_dttm: false,
         column_name: 'gender',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
       },
     ],
     expanded_columns: [
@@ -605,6 +610,7 @@ const baseQuery: QueryResponse = {
         is_dttm: true,
         column_name: 'ds',
         type: 'STRING',
+        type_generic: GenericDataType.STRING,
       },
     ],
     data: [

superset/result_set.py

@@ -29,6 +29,7 @@ from numpy.typing import NDArray
 from superset.db_engine_specs import BaseEngineSpec
 from superset.superset_typing import DbapiDescription, DbapiResult, ResultSetColumnType
 from superset.utils import core as utils
+from superset.utils.core import GenericDataType
 
 logger = logging.getLogger(__name__)
 
@@ -222,6 +223,18 @@ class SupersetResultSet:
             return False
         return column_spec.is_dttm
 
+    def type_generic(
+        self, db_type_str: Optional[str]
+    ) -> Optional[utils.GenericDataType]:
+        column_spec = self.db_engine_spec.get_column_spec(db_type_str)
+        if column_spec is None:
+            return None
+
+        if column_spec.is_dttm:
+            return GenericDataType.TEMPORAL
+
+        return column_spec.generic_type
+
     def data_type(self, col_name: str, pa_dtype: pa.DataType) -> Optional[str]:
         """Given a pyarrow data type, Returns a generic database type"""
         if set_type := self._type_dict.get(col_name):
@@ -255,7 +268,8 @@ class SupersetResultSet:
                 "column_name": col.name,
                 "name": col.name,
                 "type": db_type_str,
-                "is_dttm": self.is_temporal(db_type_str),
+                "type_generic": self.type_generic(db_type_str),
+                "is_dttm": self.is_temporal(db_type_str) or False,
             }
             columns.append(column)
         return columns

tests/integration_tests/result_set_tests.py

@@ -21,6 +21,7 @@ import tests.integration_tests.test_app
 from superset.dataframe import df_to_records
 from superset.db_engine_specs import BaseEngineSpec
 from superset.result_set import dedup, SupersetResultSet
+from superset.utils.core import GenericDataType
 
 from .base_tests import SupersetTestCase
 
@@ -48,9 +49,27 @@ class TestSupersetResultSet(SupersetTestCase):
         self.assertEqual(
             results.columns,
             [
-                {"is_dttm": False, "type": "STRING", "column_name": "a", "name": "a"},
+                {
-                {"is_dttm": False, "type": "STRING", "column_name": "b", "name": "b"},
+                    "is_dttm": False,
-                {"is_dttm": False, "type": "STRING", "column_name": "c", "name": "c"},
+                    "type": "STRING",
+                    "type_generic": GenericDataType.STRING,
+                    "column_name": "a",
+                    "name": "a",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "STRING",
+                    "type_generic": GenericDataType.STRING,
+                    "column_name": "b",
+                    "name": "b",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "STRING",
+                    "type_generic": GenericDataType.STRING,
+                    "column_name": "c",
+                    "name": "c",
+                },
             ],
         )
 
@@ -61,8 +80,20 @@ class TestSupersetResultSet(SupersetTestCase):
         self.assertEqual(
             results.columns,
             [
-                {"is_dttm": False, "type": "STRING", "column_name": "a", "name": "a"},
+                {
-                {"is_dttm": False, "type": "INT", "column_name": "b", "name": "b"},
+                    "is_dttm": False,
+                    "type": "STRING",
+                    "type_generic": GenericDataType.STRING,
+                    "column_name": "a",
+                    "name": "a",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "INT",
+                    "type_generic": GenericDataType.NUMERIC,
+                    "column_name": "b",
+                    "name": "b",
+                },
             ],
         )
 
@@ -76,11 +107,41 @@ class TestSupersetResultSet(SupersetTestCase):
         self.assertEqual(
             results.columns,
             [
-                {"is_dttm": False, "type": "FLOAT", "column_name": "a", "name": "a"},
+                {
-                {"is_dttm": False, "type": "INT", "column_name": "b", "name": "b"},
+                    "is_dttm": False,
-                {"is_dttm": False, "type": "STRING", "column_name": "c", "name": "c"},
+                    "type": "FLOAT",
-                {"is_dttm": True, "type": "DATETIME", "column_name": "d", "name": "d"},
+                    "type_generic": GenericDataType.NUMERIC,
-                {"is_dttm": False, "type": "BOOL", "column_name": "e", "name": "e"},
+                    "column_name": "a",
+                    "name": "a",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "INT",
+                    "type_generic": GenericDataType.NUMERIC,
+                    "column_name": "b",
+                    "name": "b",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "STRING",
+                    "type_generic": GenericDataType.STRING,
+                    "column_name": "c",
+                    "name": "c",
+                },
+                {
+                    "is_dttm": True,
+                    "type": "DATETIME",
+                    "type_generic": GenericDataType.TEMPORAL,
+                    "column_name": "d",
+                    "name": "d",
+                },
+                {
+                    "is_dttm": False,
+                    "type": "BOOL",
+                    "type_generic": GenericDataType.BOOLEAN,
+                    "column_name": "e",
+                    "name": "e",
+                },
             ],
         )
 
@@ -108,6 +169,7 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("user_id", "bigint", None, None, None, None, True)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "BIGINT")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.NUMERIC)
 
     def test_data_as_list_of_lists(self):
         data = [[1, "a"], [2, "b"]]
@@ -127,6 +189,7 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("is_test", "bool", None, None, None, None, True)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "BOOL")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.BOOLEAN)
         df = results.to_pandas_df()
         self.assertEqual(
             df_to_records(df),
@@ -158,9 +221,13 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("id",), ("dict_arr",), ("num_arr",), ("map_col",)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "INT")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.NUMERIC)
         self.assertEqual(results.columns[1]["type"], "STRING")
+        self.assertEqual(results.columns[1]["type_generic"], GenericDataType.STRING)
         self.assertEqual(results.columns[2]["type"], "STRING")
+        self.assertEqual(results.columns[2]["type_generic"], GenericDataType.STRING)
         self.assertEqual(results.columns[3]["type"], "STRING")
+        self.assertEqual(results.columns[3]["type_generic"], GenericDataType.STRING)
         df = results.to_pandas_df()
         self.assertEqual(
             df_to_records(df),
@@ -204,6 +271,7 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("metadata",)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "STRING")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.STRING)
         df = results.to_pandas_df()
         self.assertEqual(
             df_to_records(df),
@@ -219,6 +287,7 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("metadata",)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "STRING")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.STRING)
         df = results.to_pandas_df()
         self.assertEqual(
             df_to_records(df), [{"metadata": '[{"TestKey": [123456, "foo"]}]'}]
@@ -229,6 +298,7 @@ class TestSupersetResultSet(SupersetTestCase):
         cursor_descr = [("ds", "timestamp", None, None, None, None, True)]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "TIMESTAMP")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.TEMPORAL)
 
     def test_no_type_coercion(self):
         data = [("a", 1), ("b", 2)]
@@ -238,7 +308,9 @@ class TestSupersetResultSet(SupersetTestCase):
         ]
         results = SupersetResultSet(data, cursor_descr, BaseEngineSpec)
         self.assertEqual(results.columns[0]["type"], "VARCHAR")
+        self.assertEqual(results.columns[0]["type_generic"], GenericDataType.STRING)
         self.assertEqual(results.columns[1]["type"], "INT")
+        self.assertEqual(results.columns[1]["type_generic"], GenericDataType.NUMERIC)
 
     def test_empty_data(self):
         data = []