refactor: make DataSourceType an Enum (#6686)

Co-authored-by: Zach Nation <znation@apple.com>
Co-authored-by: GitHub Actions Bot <vega-actions-bot@users.noreply.github.com>

src/compile/data/lookup.ts

@@ -7,6 +7,7 @@ import {Model} from '../model';
 import {DataFlowNode, OutputNode} from './dataflow';
 import {findSource} from './parse';
 import {SourceNode} from './source';
+import {DataSourceType} from '../../data';
 
 export class LookupNode extends DataFlowNode {
   public clone() {
@@ -31,7 +32,12 @@ export class LookupNode extends DataFlowNode {
       }
 
       const fromOutputName = model.getName(`lookup_${counter}`);
-      fromOutputNode = new OutputNode(fromSource, fromOutputName, 'lookup', model.component.data.outputNodeRefCounts);
+      fromOutputNode = new OutputNode(
+        fromSource,
+        fromOutputName,
+        DataSourceType.Lookup,
+        model.component.data.outputNodeRefCounts
+      );
       model.component.data.outputNodes[fromOutputName] = fromOutputNode;
     } else if (isLookupSelection(from)) {
       const selName = from.selection;

src/compile/data/optimizers.ts

@@ -1,4 +1,4 @@
-import {MAIN, Parse} from '../../data';
+import {DataSourceType, Parse} from '../../data';
 import {Dict, fieldIntersection, hash, hasIntersection, isEmpty, keys, some} from '../../util';
 import {Model} from '../model';
 import {requiresSelectionId} from '../selection';
@@ -244,7 +244,7 @@ export function moveFacetDown(node: DataFlowNode) {
 }
 
 function moveMainDownToFacet(node: DataFlowNode) {
-  if (node instanceof OutputNode && node.type === MAIN) {
+  if (node instanceof OutputNode && node.type === DataSourceType.Main) {
     if (node.numChildren() === 1) {
       const child = node.children[0];
       if (!(child instanceof FacetNode)) {

src/compile/data/parse.ts

@@ -7,9 +7,8 @@ import {
   isNamedData,
   isSequenceGenerator,
   isUrlData,
-  MAIN,
-  ParseValue,
-  RAW
+  DataSourceType,
+  ParseValue
 } from '../../data';
 import * as log from '../../log';
 import {
@@ -359,8 +358,8 @@ export function parseData(model: Model): DataComponent {
   }
 
   // add an output node pre aggregation
-  const rawName = model.getName(RAW);
-  const raw = new OutputNode(head, rawName, RAW, outputNodeRefCounts);
+  const rawName = model.getDataName(DataSourceType.Raw);
+  const raw = new OutputNode(head, rawName, DataSourceType.Raw, outputNodeRefCounts);
   outputNodes[rawName] = raw;
   head = raw;
 
@@ -382,8 +381,8 @@ export function parseData(model: Model): DataComponent {
   }
 
   // output node for marks
-  const mainName = model.getName(MAIN);
-  const main = new OutputNode(head, mainName, MAIN, outputNodeRefCounts);
+  const mainName = model.getDataName(DataSourceType.Main);
+  const main = new OutputNode(head, mainName, DataSourceType.Main, outputNodeRefCounts);
   outputNodes[mainName] = main;
   head = main;
 

src/compile/mark/mark.ts

@@ -1,6 +1,6 @@
 import {isArray} from 'vega-util';
 import {FieldRefOption, isFieldDef, isValueDef, vgField} from '../../channeldef';
-import {MAIN} from '../../data';
+import {DataSourceType} from '../../data';
 import {isAggregate, pathGroupingFields} from '../../encoding';
 import {AREA, BAR, isPathMark, LINE, Mark, TRAIL} from '../../mark';
 import {isSortByEncoding, isSortField} from '../../sort';
@@ -68,8 +68,8 @@ function getPathGroups(model: UnitModel, details: string[]) {
       type: 'group',
       from: {
         facet: {
-          name: FACETED_PATH_PREFIX + model.requestDataName(MAIN),
-          data: model.requestDataName(MAIN),
+          name: FACETED_PATH_PREFIX + model.requestDataName(DataSourceType.Main),
+          data: model.requestDataName(DataSourceType.Main),
           groupby: details
         }
       },
@@ -211,8 +211,8 @@ function getGroupsForStackedBarWithCornerRadius(model: UnitModel) {
       type: 'group',
       from: {
         facet: {
-          data: model.requestDataName(MAIN),
-          name: STACK_GROUP_PREFIX + model.requestDataName(MAIN),
+          data: model.requestDataName(DataSourceType.Main),
+          name: STACK_GROUP_PREFIX + model.requestDataName(DataSourceType.Main),
           groupby,
           aggregate: {
             fields: [
@@ -320,7 +320,7 @@ function getMarkGroup(model: UnitModel, opt: {fromPrefix: string} = {fromPrefix:
       ...(sort ? {sort} : {}),
       ...(interactive ? interactive : {}),
       ...(aria === false ? {aria} : {}),
-      from: {data: opt.fromPrefix + model.requestDataName(MAIN)},
+      from: {data: opt.fromPrefix + model.requestDataName(DataSourceType.Main)},
       encode: {
         update: markCompiler[mark].encodeEntry(model)
       },

src/compile/model.ts

@@ -469,13 +469,17 @@ export abstract class Model {
     return varName((this.name ? this.name + '_' : '') + text);
   }
 
+  public getDataName(type: DataSourceType) {
+    return this.getName(DataSourceType[type].toLowerCase());
+  }
+
   /**
    * Request a data source name for the given data source type and mark that data source as required.
    * This method should be called in parse, so that all used data source can be correctly instantiated in assembleData().
    * You can lookup the correct dataset name in assemble with `lookupDataSource`.
    */
   public requestDataName(name: DataSourceType) {
-    const fullName = this.getName(name);
+    const fullName = this.getDataName(name);
 
     // Increase ref count. This is critical because otherwise we won't create a data source.
     // We also increase the ref counts on OutputNode.getSource() calls.

src/compile/projection/parse.ts

@@ -2,7 +2,7 @@ import {SignalRef} from 'vega';
 import {hasOwnProperty} from 'vega-util';
 import {LATITUDE, LATITUDE2, LONGITUDE, LONGITUDE2, SHAPE} from '../../channel';
 import {getFieldOrDatumDef} from '../../channeldef';
-import {MAIN} from '../../data';
+import {DataSourceType} from '../../data';
 import {PROJECTION_PROPERTIES} from '../../projection';
 import {GEOJSON} from '../../type';
 import {duplicate, every, stringify} from '../../util';
@@ -59,7 +59,7 @@ function gatherFitData(model: UnitModel) {
 
   if (data.length === 0) {
     // main source is geojson, so we can just use that
-    data.push(model.requestDataName(MAIN));
+    data.push(model.requestDataName(DataSourceType.Main));
   }
 
   return data;

src/compile/scale/domain.ts

@@ -22,7 +22,7 @@ import {
   valueExpr,
   vgField
 } from '../../channeldef';
-import {MAIN, RAW} from '../../data';
+import {DataSourceType} from '../../data';
 import {DateTime} from '../../datetime';
 import * as log from '../../log';
 import {Domain, hasDiscreteDomain, isDomainUnionWith, isSelectionDomain, ScaleConfig, ScaleType} from '../../scale';
@@ -261,7 +261,7 @@ function parseSingleChannelDomain(
       return makeImplicit([[0, 1]]);
     }
 
-    const data = model.requestDataName(MAIN);
+    const data = model.requestDataName(DataSourceType.Main);
     return makeImplicit([
       {
         data,
@@ -284,7 +284,7 @@ function parseSingleChannelDomain(
 
   const fieldDef = fieldOrDatumDef; // now we can be sure it's a fieldDef
   if (domain === 'unaggregated') {
-    const data = model.requestDataName(MAIN);
+    const data = model.requestDataName(DataSourceType.Main);
     const {field} = fieldOrDatumDef;
     return makeImplicit([
       {
@@ -309,7 +309,9 @@ function parseSingleChannelDomain(
         {
           // If sort by aggregation of a specified sort field, we need to use RAW table,
           // so we can aggregate values for the scale independently from the main aggregation.
-          data: util.isBoolean(sort) ? model.requestDataName(MAIN) : model.requestDataName(RAW),
+          data: util.isBoolean(sort)
+            ? model.requestDataName(DataSourceType.Main)
+            : model.requestDataName(DataSourceType.Raw),
           // Use range if we added it and the scale does not support computing a range as a signal.
           field: model.vgField(channel, binRequiresRange(fieldDef, channel) ? {binSuffix: 'range'} : {}),
           // we have to use a sort object if sort = true to make the sort correct by bin start
@@ -336,7 +338,7 @@ function parseSingleChannelDomain(
       } else {
         return makeImplicit([
           {
-            data: model.requestDataName(MAIN),
+            data: model.requestDataName(DataSourceType.Main),
             field: model.vgField(channel, {})
           }
         ]);
@@ -354,7 +356,7 @@ function parseSingleChannelDomain(
       model.config
     )
   ) {
-    const data = model.requestDataName(MAIN);
+    const data = model.requestDataName(DataSourceType.Main);
     return makeImplicit([
       {
         data,
@@ -370,15 +372,17 @@ function parseSingleChannelDomain(
       {
         // If sort by aggregation of a specified sort field, we need to use RAW table,
         // so we can aggregate values for the scale independently from the main aggregation.
-        data: util.isBoolean(sort) ? model.requestDataName(MAIN) : model.requestDataName(RAW),
+        data: util.isBoolean(sort)
+          ? model.requestDataName(DataSourceType.Main)
+          : model.requestDataName(DataSourceType.Raw),
         field: model.vgField(channel),
         sort: sort
       }
     ]);
   } else {
     return makeImplicit([
       {
-        data: model.requestDataName(MAIN),
+        data: model.requestDataName(DataSourceType.Main),
         field: model.vgField(channel)
       }
     ]);

src/compile/scale/range.ts

@@ -24,7 +24,7 @@ import {
 } from '../../channel';
 import {getFieldOrDatumDef, ScaleDatumDef, ScaleFieldDef} from '../../channeldef';
 import {Config, getViewConfigDiscreteSize, getViewConfigDiscreteStep, ViewConfig} from '../../config';
-import {MAIN} from '../../data';
+import {DataSourceType} from '../../data';
 import * as log from '../../log';
 import {Mark} from '../../mark';
 import {
@@ -132,7 +132,7 @@ export function parseRangeForChannel(channel: ScaleChannel, model: UnitModel): E
               }
             } else if (isObject(range)) {
               return makeExplicit({
-                data: model.requestDataName(MAIN),
+                data: model.requestDataName(DataSourceType.Main),
                 field: range.field,
                 sort: {op: 'min', field: model.vgField(channel)}
               });

src/compile/selection/parse.ts

@@ -10,6 +10,7 @@ import {FilterNode} from '../data/filter';
 import {Model} from '../model';
 import {UnitModel} from '../unit';
 import {forEachTransform} from './transforms/transforms';
+import {DataSourceType} from '../../data';
 
 export function parseUnitSelection(model: UnitModel, selDefs: Dict<SelectionDef>) {
   const selCmpts: Dict<SelectionComponent<any /* this has to be "any" so typing won't fail in test files*/>> = {};
@@ -129,7 +130,7 @@ export function materializeSelections(model: UnitModel, main: OutputNode) {
     model.component.data.outputNodes[lookupName] = selCmpt.materialized = new OutputNode(
       new FilterNode(main, model, {selection}),
       lookupName,
-      'lookup',
+      DataSourceType.Lookup,
       model.component.data.outputNodeRefCounts
     );
   });

src/data.ts

@@ -148,10 +148,13 @@ export function isGraticuleGenerator(data: Partial<Data> | Partial<VgData>): dat
   return 'graticule' in data;
 }
 
-export type DataSourceType = 'raw' | 'main' | 'row' | 'column' | 'lookup';
-
-export const MAIN = 'main' as const;
-export const RAW = 'raw' as const;
+export enum DataSourceType {
+  Raw,
+  Main,
+  Row,
+  Column,
+  Lookup
+}
 
 export type Generator = SequenceGenerator | SphereGenerator | GraticuleGenerator;
 

test/compile/data/assemble.test.ts

@@ -4,13 +4,14 @@ import {OutputNode} from '../../../src/compile/data/dataflow';
 import {SourceNode} from '../../../src/compile/data/source';
 import {WindowTransformNode} from '../../../src/compile/data/window';
 import {Transform} from '../../../src/transform';
+import {DataSourceType} from '../../../src/data';
 
 describe('compile/data/assemble', () => {
   describe('assembleData', () => {
     it('should assemble named data source', () => {
       const src = new SourceNode({name: 'foo'});
       const outputNodeRefCounts = {};
-      const main = new OutputNode(null, 'mainOut', 'main', outputNodeRefCounts);
+      const main = new OutputNode(null, 'mainOut', DataSourceType.Main, outputNodeRefCounts);
       main.parent = src;
 
       expect(main.getSource()).toBe('mainOut');
@@ -32,11 +33,11 @@ describe('compile/data/assemble', () => {
     it('should assemble raw and main output', () => {
       const src = new SourceNode({url: 'foo.csv'});
       const outputNodeRefCounts = {};
-      const raw = new OutputNode(null, 'rawOut', 'raw', outputNodeRefCounts);
+      const raw = new OutputNode(null, 'rawOut', DataSourceType.Raw, outputNodeRefCounts);
       raw.parent = src;
       const agg = new AggregateNode(null, new Set(['a']), {b: {count: new Set(['count_*'])}});
       agg.parent = raw;
-      const main = new OutputNode(null, 'mainOut', 'main', outputNodeRefCounts);
+      const main = new OutputNode(null, 'mainOut', DataSourceType.Main, outputNodeRefCounts);
       main.parent = agg;
 
       expect(raw.getSource()).toBe('rawOut');
@@ -77,7 +78,7 @@ describe('compile/data/assemble', () => {
     it('should assemble window transform node', () => {
       const src = new SourceNode({url: 'foo.csv'});
       const outputNodeRefCounts = {};
-      const raw = new OutputNode(null, 'rawOut', 'raw', outputNodeRefCounts);
+      const raw = new OutputNode(null, 'rawOut', DataSourceType.Raw, outputNodeRefCounts);
       raw.parent = src;
       const transform: Transform = {
         window: [
@@ -98,7 +99,7 @@ describe('compile/data/assemble', () => {
       };
       const agg = new WindowTransformNode(null, transform);
       agg.parent = raw;
-      const main = new OutputNode(null, 'mainOut', 'main', outputNodeRefCounts);
+      const main = new OutputNode(null, 'mainOut', DataSourceType.Main, outputNodeRefCounts);
       main.parent = agg;
 
       expect(raw.getSource()).toBe('rawOut');
@@ -146,7 +147,7 @@ describe('compile/data/assemble', () => {
     it('should assemble named datasets with datastore', () => {
       const src = new SourceNode({name: 'foo'});
       const outputNodeRefCounts = {};
-      const main = new OutputNode(null, 'mainOut', 'main', outputNodeRefCounts);
+      const main = new OutputNode(null, 'mainOut', DataSourceType.Main, outputNodeRefCounts);
       main.parent = src;
 
       const data = assembleRootData(

test/compile/data/dataflow.test.ts

@@ -1,4 +1,5 @@
 import {OutputNode} from '../../../src/compile/data/dataflow';
+import {DataSourceType} from '../../../src/data';
 import * as log from '../../../src/log';
 import {PlaceholderDataFlowNode} from './util';
 
@@ -161,14 +162,14 @@ describe('compile/data/dataflow', () => {
   describe('OutputNode', () => {
     describe('dependentFields', () => {
       it('should return empty set', () => {
-        const flatten = new OutputNode(null, 'src', 'main', {});
+        const flatten = new OutputNode(null, 'src', DataSourceType.Main, {});
         expect(flatten.dependentFields()).toEqual(new Set());
       });
     });
 
     describe('producedFields', () => {
       it('should return empty set', () => {
-        const flatten = new OutputNode(null, 'src', 'main', {});
+        const flatten = new OutputNode(null, 'src', DataSourceType.Main, {});
         expect(flatten.producedFields()).toEqual(new Set());
       });
     });

test/compile/model.test.ts

@@ -1,5 +1,6 @@
 import {NameMap} from '../../src/compile/model';
 import {parseFacetModelWithScale, parseModel} from '../util';
+import {DataSourceType} from '../../src/data';
 
 describe('Model', () => {
   describe('NameMap', () => {
@@ -45,6 +46,20 @@ describe('Model', () => {
     });
   });
 
+  describe('getDataName', () => {
+    it('returns correct names for DataSourceType', () => {
+      const model = parseModel({
+        data: {values: []},
+        mark: 'point'
+      });
+      expect(model.getDataName(DataSourceType.Column)).toBe('column');
+      expect(model.getDataName(DataSourceType.Lookup)).toBe('lookup');
+      expect(model.getDataName(DataSourceType.Main)).toBe('main');
+      expect(model.getDataName(DataSourceType.Raw)).toBe('raw');
+      expect(model.getDataName(DataSourceType.Row)).toBe('row');
+    });
+  });
+
   describe('getSizeSignalRef', () => {
     it('returns formula for step if parent is facet', () => {
       const model = parseFacetModelWithScale({