Merge pull request #883 from concord-consortium/186033205-formula-count

feat(formula): add `count` function [PT-186033205]

v3/src/models/data/formula-fn-registry.ts

@@ -148,6 +148,24 @@ export const fnRegistry = {
     }
   },
 
+  // count(expression, filterExpression)
+  count: {
+    isAggregate: true,
+    cachedEvaluateFactory: cachedAggregateFnFactory,
+    evaluateRaw: (args: MathNode[], mathjs: any, scope: FormulaMathJsScope) => {
+      const expression = args[0]
+      const filter = args[1]
+      if (!expression) {
+        // Special case - count() without arguments returns number of children cases.
+        return scope.getCaseChildrenCount()
+      }
+      let expressionValues = evaluateNode(expression, scope)
+      const filterValues = filter !== undefined ? evaluateNode(filter, scope) : undefined
+      expressionValues = expressionValues.filter((v: any, i: number) => v !== "" && (filter ? !!filterValues[i] : true))
+      return expressionValues.length
+    }
+  },
+
   // next(expression, defaultValue, filter)
   next: {
     // expression and filter are evaluated as aggregate symbols, defaultValue is not - it depends on case index

v3/src/models/data/formula-manager.ts

@@ -105,6 +105,24 @@ export class FormulaManager {
     return caseGroupId
   }
 
+  getCaseChildrenCountMap(formulaId: string) {
+    const { attributeId, dataSet } = this.getFormulaContext(formulaId)
+
+    const collectionId = dataSet.getCollectionForAttribute(attributeId)?.id
+    const collectionIndex = dataSet.getCollectionIndex(collectionId || "")
+    const caseChildrenCount: Record<string, number> = {}
+
+    const formulaCollection = dataSet.collectionGroups[collectionIndex]
+    if (formulaCollection) {
+      dataSet.collectionGroups[collectionIndex].groups.forEach((group: CaseGroup) =>
+        caseChildrenCount[group.pseudoCase.__id__] =
+          group.childPseudoCaseIds ? group.childPseudoCaseIds.length : group.childCaseIds.length
+      )
+    }
+
+    return caseChildrenCount
+  }
+
   recalculateFormula(formulaId: string, casesToRecalculateDesc: ICase[] | "ALL_CASES" = "ALL_CASES") {
     const { formula, attributeId, dataSet } = this.getFormulaContext(formulaId)
 
@@ -154,7 +172,8 @@ export class FormulaManager {
       //   referencing attributes from child collections.
       useSameLevelGrouping: collectionIndex === childMostAggregateCollectionIndex,
       childMostCollectionCases,
-      caseGroupId: this.getCaseGroupMap(formulaId)
+      caseGroupId: this.getCaseGroupMap(formulaId),
+      caseChildrenCount: this.getCaseChildrenCountMap(formulaId)
     })
 
     let compiledFormula: EvalFunction

v3/src/models/data/formula-mathjs-scope.ts

@@ -12,6 +12,7 @@ export interface IFormulaMathjsScopeContext {
   childMostCollectionCases: ICase[]
   useSameLevelGrouping: boolean
   caseGroupId: Record<string, string>
+  caseChildrenCount: Record<string, number>
   globalValueManager?: IGlobalValueManager
 }
 
@@ -120,6 +121,10 @@ export class FormulaMathJsScope {
     return this.context.localDataSet.caseIDMap[this.caseId]
   }
 
+  getCaseChildrenCount() {
+    return this.context.caseChildrenCount[this.caseId] ?? 0
+  }
+
   getLocalDataSet() {
     return this.context.localDataSet
   }