Merge pull request #935 from concord-consortium/186201910-fix-lookup-fn

fix(formula): fix lookup functions [PT-186201910]
concord-consortium · Oct 11, 2023 · 3c95212 · 3c95212
2 parents de15b17 + 62aa587
commit 3c95212
Showing 1 changed file with 9 additions and 13 deletions.
diff --git a/v3/src/models/data/formula-fn-registry.ts b/v3/src/models/data/formula-fn-registry.ts
@@ -112,16 +112,16 @@ export const fnRegistry = {
       }
       if (!isConstantStringNode(args[0]) || !isConstantStringNode(args[1])) {
         throw new Error("lookupByIndex function expects first two arguments to be strings " +
-          "and the third one to be numeric")
+          "and the third one to be numeric (or evaluate to number)")
       }
       return [args[0], args[1], args[2]]
     },
     getDependency: (args: MathNode[]): ILookupDependency => {
       const validArgs = fnRegistry.lookupByIndex.validateArguments(args)
       return {
         type: "lookup",
-        dataSetId: validArgs[0].value,
-        attrId: validArgs[1].value,
+        dataSetId: rmCanonicalPrefix(validArgs[0].value),
+        attrId: rmCanonicalPrefix(validArgs[1].value),
       }
     },
     canonicalize: (args: MathNode[], displayNameMap: DisplayNameMap) => {
@@ -132,8 +132,7 @@ export const fnRegistry = {
       validArgs[1].value = displayNameMap.dataSet[dataSetName]?.attribute[attrName]
     },
     evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
-      const dataSetId = rmCanonicalPrefix(evaluateNode(args[0], scope))
-      const attrId = rmCanonicalPrefix(evaluateNode(args[1], scope))
+      const { dataSetId, attrId } = fnRegistry.lookupByIndex.getDependency(args)
       const zeroBasedIndex = evaluateNode(args[2], scope) - 1
       return scope.getDataSet(dataSetId)?.getValueAtIndex(zeroBasedIndex, attrId) || UNDEF_RESULT
     }
@@ -151,13 +150,13 @@ export const fnRegistry = {
       }
       return [args[0], args[1], args[2], args[3]]
     },
-    getDependency: (args: MathNode[]): ILookupDependency => {
+    getDependency: (args: MathNode[]): Required<ILookupDependency> => {
       const validArgs = fnRegistry.lookupByKey.validateArguments(args)
       return {
         type: "lookup",
-        dataSetId: validArgs[0].value,
-        attrId: validArgs[1].value,
-        keyAttrId: validArgs[2].value,
+        dataSetId: rmCanonicalPrefix(validArgs[0].value),
+        attrId: rmCanonicalPrefix(validArgs[1].value),
+        keyAttrId: rmCanonicalPrefix(validArgs[2].value),
       }
     },
     canonicalize: (args: MathNode[], displayNameMap: DisplayNameMap) => {
@@ -170,11 +169,8 @@ export const fnRegistry = {
       validArgs[2].value = displayNameMap.dataSet[dataSetName]?.attribute[keyAttrName]
     },
     evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
-      const dataSetId = rmCanonicalPrefix(evaluateNode(args[0], scope))
-      const attrId = rmCanonicalPrefix(evaluateNode(args[1], scope))
-      const keyAttrId = rmCanonicalPrefix(evaluateNode(args[2], scope))
+      const { dataSetId, attrId, keyAttrId } = fnRegistry.lookupByKey.getDependency(args)
       const keyAttrValue = evaluateNode(args[3], scope)
-
       const dataSet: IDataSet | undefined = scope.getDataSet(dataSetId)
       if (!dataSet) {
         return UNDEF_RESULT