@tanstack/db/dist/esm/query/compiler/group-by.js

A reactive client store for building super fast apps on sync

import { groupBy, map, serializeValue, filter, groupByOperators } from "@tanstack/db-ivm";
import { isExpressionLike, getHavingExpression, ConditionalSelect, PropRef, Func } from "../ir.js";
import { NonAggregateExpressionNotInGroupByError, UnsupportedAggregateFunctionError, UnknownHavingExpressionTypeError, AggregateFunctionNotInSelectError } from "../../errors.js";
import { getParentContextIdentity, getEqualityValueIdentity, getParentContextValue } from "../equality-value-identity.js";
import { compileExpression, toBooleanPredicate, isCaseWhenConditionTrue } from "./evaluators.js";
import { attachRouteMetadata, stripInternalCallbackMetadata, getNamespacedRouteMetadata, INCLUDES_PUBLIC_KEY } from "./route-metadata.js";
const RAW_REPRESENTATIVE = /* @__PURE__ */ Symbol(`raw_group_representative`);
function createInternalGroupFields(groupCount, selectClause) {
  const aliases = Object.keys(selectClause ?? {});
  let prefix = `__tanstack_group_`;
  while (aliases.some((alias) => alias.startsWith(prefix))) prefix += `_`;
  return {
    virtual: `${prefix}virtual`,
    route: `${prefix}route`,
    correlationIdentity: `${prefix}correlation_identity`,
    parentContextIdentity: `${prefix}parent_context_identity`,
    singleGroup: `${prefix}single_group`,
    aggregatePrefix: `${prefix}aggregate_`,
    groupKeys: Array.from(
      { length: groupCount },
      (_, i) => `${prefix}key_${i}`
    ),
    groupValues: Array.from(
      { length: groupCount },
      (_, i) => `${prefix}value_${i}`
    ),
    groupKeyRefs: Array.from(
      { length: groupCount },
      (_, i) => `${prefix}key_ref_${i}`
    )
  };
}
function createPublicGroupKey(values) {
  const identities = values.map(getEqualityValueIdentity);
  if (identities.length === 1) {
    const identity = identities[0];
    if (identity == null || typeof identity !== `object` && typeof identity !== `function` && typeof identity !== `symbol`) {
      return identity;
    }
  }
  return serializeValue(identities);
}
function attachPublicGroupKey(row, publicKey) {
  const keyedRow = row;
  keyedRow[INCLUDES_PUBLIC_KEY] = publicKey;
}
function createRepresentative(rowKey, value, identity) {
  const representative = {
    key: serializeValue([rowKey, identity])
  };
  Object.defineProperty(representative, RAW_REPRESENTATIVE, { value });
  return representative;
}
function getRepresentative(values) {
  let selected;
  for (const [candidate, multiplicity] of values) {
    if (multiplicity <= 0) continue;
    if (selected === void 0 || candidate.key < selected.key) {
      selected = candidate;
    }
  }
  return selected;
}
function unwrapRepresentative(value) {
  return value?.[RAW_REPRESENTATIVE];
}
function addCorrelationRouteIdentityToGroupKey(key, row, mainSource, fields, valueIdentity) {
  const route = getNamespacedRouteMetadata(row, mainSource);
  key[fields.correlationIdentity] = valueIdentity.equality(
    route?.correlationKey
  );
  if (route?.parentContext != null) {
    key[fields.parentContextIdentity] = getParentContextIdentity(
      route.parentContext
    );
  }
}
function addCorrelationRouteAggregate(aggregates, mainSource, fields, valueIdentity) {
  aggregates[fields.route] = {
    preMap: ([rowKey, row]) => {
      const route = getNamespacedRouteMetadata(row, mainSource);
      return createRepresentative(rowKey, route, [
        valueIdentity.exact(route?.correlationKey),
        getParentContextIdentity(route?.parentContext)
      ]);
    },
    reduce: getRepresentative,
    postMap: unwrapRepresentative
  };
}
function getGroupRoute(aggregatedRow, fields) {
  return aggregatedRow[fields.route];
}
function getCorrelationRouteIdentity(aggregatedRow, fields) {
  return getGroupRoute(aggregatedRow, fields)?.parentContext == null ? aggregatedRow[fields.correlationIdentity] : [
    aggregatedRow[fields.correlationIdentity],
    aggregatedRow[fields.parentContextIdentity]
  ];
}
function getGroupEvaluationRow(row, fields, selected = row.$selected) {
  return {
    ...getParentContextValue(getGroupRoute(row, fields)?.parentContext),
    $selected: selected
  };
}
function getRowVirtualMetadata(row) {
  let found = false;
  let allSynced = true;
  let hasLocal = false;
  for (const [alias, value] of Object.entries(row)) {
    if (alias === `$selected`) continue;
    if (value === null || typeof value !== `object`) continue;
    const asRecord = value;
    const hasSyncedProp = `$synced` in asRecord;
    const hasOriginProp = `$origin` in asRecord;
    if (!hasSyncedProp && !hasOriginProp) {
      continue;
    }
    found = true;
    if (asRecord.$synced === false) {
      allSynced = false;
    }
    if (asRecord.$origin === `local`) {
      hasLocal = true;
    }
  }
  return {
    synced: found ? allSynced : true,
    hasLocal
  };
}
const { sum, count, avg, min, max } = groupByOperators;
function validateAndCreateMapping(groupByClause, selectClause) {
  const selectToGroupByIndex = /* @__PURE__ */ new Map();
  if (!selectClause) {
    return selectToGroupByIndex;
  }
  for (const [alias, expr] of Object.entries(selectClause)) {
    if (expr.type === `agg` || containsAggregate(expr)) {
      continue;
    }
    const groupIndex = groupByClause.findIndex(
      (groupExpr) => expressionsEqual(expr, groupExpr)
    );
    if (groupIndex === -1) {
      throw new NonAggregateExpressionNotInGroupByError(alias);
    }
    selectToGroupByIndex.set(alias, groupIndex);
  }
  return selectToGroupByIndex;
}
function processGroupBy(pipeline, groupByClause, valueIdentity, havingClauses, selectClause, fnHavingClauses, aggregateCollectionId, mainSource, sanitizeCallbackRows = false) {
  const fields = createInternalGroupFields(groupByClause.length, selectClause);
  const virtualAggregates = {
    [fields.virtual]: {
      preMap: ([, row]) => getRowVirtualMetadata(row),
      reduce: (values) => {
        const group = { synced: true, hasLocal: false };
        for (const [metadata, multiplicity] of values) {
          if (multiplicity <= 0) continue;
          if (!metadata.synced) group.synced = false;
          if (metadata.hasLocal) group.hasLocal = true;
        }
        return group;
      }
    }
  };
  if (mainSource) {
    addCorrelationRouteAggregate(
      virtualAggregates,
      mainSource,
      fields,
      valueIdentity
    );
  }
  const singleGroup = groupByClause.length === 0;
  const mapping = singleGroup ? void 0 : validateAndCreateMapping(groupByClause, selectClause);
  const compiledGroupByExpressions = groupByClause.map(
    (e) => compileExpression(e)
  );
  const keyExtractor = ([, row]) => {
    const namespacedRow = singleGroup ? row : { ...row };
    if (!singleGroup) delete namespacedRow.$selected;
    const key = singleGroup ? { [fields.singleGroup]: true } : {};
    for (let i = 0; i < groupByClause.length; i++) {
      const compiledExpr = compiledGroupByExpressions[i];
      const value = compiledExpr(namespacedRow);
      key[fields.groupKeys[i]] = valueIdentity.equality(value);
    }
    if (mainSource) {
      addCorrelationRouteIdentityToGroupKey(
        key,
        row,
        mainSource,
        fields,
        valueIdentity
      );
    }
    return key;
  };
  const aggregates = virtualAggregates;
  const wrappedAggExprs = {};
  const aggCounter = { value: 0 };
  for (let i = 0; i < compiledGroupByExpressions.length; i++) {
    const compiledExpr = compiledGroupByExpressions[i];
    aggregates[fields.groupValues[i]] = {
      preMap: ([rowKey, row]) => {
        const value = compiledExpr(row);
        return createRepresentative(rowKey, value, valueIdentity.exact(value));
      },
      reduce: getRepresentative,
      postMap: unwrapRepresentative
    };
  }
  if (selectClause) {
    for (const [alias, expr] of Object.entries(selectClause)) {
      if (expr.type === `agg`) {
        aggregates[alias] = getAggregateFunction(expr);
      } else if (containsAggregate(expr)) {
        const { transformed, extracted } = extractAndReplaceAggregates(
          expr,
          aggCounter,
          fields.aggregatePrefix
        );
        for (const [syntheticAlias, aggExpr] of Object.entries(extracted)) {
          aggregates[syntheticAlias] = getAggregateFunction(aggExpr);
        }
        wrappedAggExprs[alias] = compileGroupedSelectValue(
          singleGroup ? transformed : replaceGroupByRefsInSelectValue(
            transformed,
            groupByClause,
            fields.groupKeyRefs
          )
        );
      }
    }
  }
  pipeline = pipeline.pipe(groupBy(keyExtractor, aggregates));
  pipeline = pipeline.pipe(
    map(([, aggregatedRow]) => {
      const selectResults = aggregatedRow.$selected || {};
      const finalResults = singleGroup ? { ...selectResults } : {};
      if (selectClause) {
        for (const [alias, expr] of Object.entries(selectClause)) {
          if (expr.type === `agg`) {
            finalResults[alias] = aggregatedRow[alias];
          } else if (!singleGroup && !wrappedAggExprs[alias]) {
            const groupIndex = mapping?.get(alias);
            if (groupIndex !== void 0) {
              finalResults[alias] = aggregatedRow[fields.groupValues[groupIndex]];
            } else {
              finalResults[alias] = selectResults[alias];
            }
          }
        }
        evaluateWrappedAggregates(
          finalResults,
          aggregatedRow,
          wrappedAggExprs,
          fields
        );
      } else {
        for (let i = 0; i < groupByClause.length; i++) {
          finalResults[`__key_${i}`] = aggregatedRow[fields.groupValues[i]];
        }
      }
      const route = mainSource ? getGroupRoute(aggregatedRow, fields) : void 0;
      const correlationKey = route?.correlationKey;
      const correlationRoute = mainSource ? getCorrelationRouteIdentity(aggregatedRow, fields) : void 0;
      const keyParts = [];
      const publicKeyParts = [];
      for (let i = 0; i < groupByClause.length; i++) {
        keyParts.push(aggregatedRow[fields.groupKeys[i]]);
        publicKeyParts.push(aggregatedRow[fields.groupValues[i]]);
      }
      if (correlationRoute !== void 0) {
        keyParts.push(correlationRoute);
      }
      const finalKey = singleGroup ? correlationRoute !== void 0 ? `single_group_${serializeValue(correlationRoute)}` : `single_group` : keyParts.length === 1 ? keyParts[0] : serializeValue(keyParts);
      const publicKey = singleGroup ? `single_group` : createPublicGroupKey(publicKeyParts);
      const resultRow = {
        ...aggregatedRow,
        $selected: finalResults
      };
      const virtual = aggregatedRow[fields.virtual];
      resultRow.$synced = virtual?.synced ?? true;
      resultRow.$origin = virtual?.hasLocal ? `local` : `remote`;
      resultRow.$key = publicKey;
      resultRow.$collectionId = aggregateCollectionId ?? resultRow.$collectionId;
      if (mainSource && correlationKey !== void 0) {
        attachPublicGroupKey(resultRow, publicKey);
        attachRouteMetadata(
          resultRow,
          correlationKey,
          route?.parentContext ?? null
        );
      }
      return [mainSource ? finalKey : publicKey, resultRow];
    })
  );
  if (havingClauses && havingClauses.length > 0) {
    for (const havingClause of havingClauses) {
      const havingExpression = getHavingExpression(havingClause);
      const transformedHavingClause = replaceAggregatesByRefs(
        havingExpression,
        selectClause || {}
      );
      const compiledHaving = compileExpression(transformedHavingClause);
      pipeline = pipeline.pipe(
        filter(([, row]) => {
          const namespacedRow = getGroupEvaluationRow(row, fields);
          const result = compiledHaving(namespacedRow);
          return singleGroup ? toBooleanPredicate(result) : result;
        })
      );
    }
  }
  if (fnHavingClauses && fnHavingClauses.length > 0) {
    for (const fnHaving of fnHavingClauses) {
      pipeline = pipeline.pipe(
        filter(([, row]) => {
          const namespacedRow = getGroupEvaluationRow(row, fields);
          const callbackRow = sanitizeCallbackRows ? stripInternalCallbackMetadata(namespacedRow) : namespacedRow;
          return toBooleanPredicate(fnHaving(callbackRow));
        })
      );
    }
  }
  return pipeline;
}
function expressionsEqual(expr1, expr2) {
  if (!expr1 || !expr2) return false;
  if (expr1.type !== expr2.type) return false;
  switch (expr1.type) {
    case `ref`:
      if (!expr1.path || !expr2.path) return false;
      if (expr1.path.length !== expr2.path.length) return false;
      return expr1.path.every(
        (segment, i) => segment === expr2.path[i]
      );
    case `val`:
      return expr1.value === expr2.value;
    case `func`:
      return expr1.name === expr2.name && expr1.args?.length === expr2.args?.length && (expr1.args || []).every(
        (arg, i) => expressionsEqual(arg, expr2.args[i])
      );
    case `agg`:
      return expr1.name === expr2.name && expr1.args?.length === expr2.args?.length && (expr1.args || []).every(
        (arg, i) => expressionsEqual(arg, expr2.args[i])
      );
    default:
      return false;
  }
}
function getAggregateFunction(aggExpr) {
  const compiledExpr = compileExpression(aggExpr.args[0]);
  const valueExtractor = ([, namespacedRow]) => {
    const value = compiledExpr(namespacedRow);
    if (typeof value === `number`) {
      return value;
    }
    return value != null ? Number(value) : 0;
  };
  const valueExtractorForMinMax = ([, namespacedRow]) => {
    const value = compiledExpr(namespacedRow);
    if (typeof value === `number` || typeof value === `string` || typeof value === `bigint` || value instanceof Date) {
      return value;
    }
    return value != null ? Number(value) : 0;
  };
  const rawValueExtractor = ([, namespacedRow]) => {
    return compiledExpr(namespacedRow);
  };
  switch (aggExpr.name.toLowerCase()) {
    case `sum`:
      return sum(valueExtractor);
    case `count`:
      return count(rawValueExtractor);
    case `avg`:
      return avg(valueExtractor);
    case `min`:
      return min(valueExtractorForMinMax);
    case `max`:
      return max(valueExtractorForMinMax);
    default:
      throw new UnsupportedAggregateFunctionError(aggExpr.name);
  }
}
function replaceAggregatesByRefs(havingExpr, selectClause, resultAlias = `$selected`) {
  switch (havingExpr.type) {
    case `agg`: {
      const aggExpr = havingExpr;
      for (const [alias, selectExpr] of Object.entries(selectClause)) {
        if (selectExpr.type === `agg` && aggregatesEqual(aggExpr, selectExpr)) {
          return new PropRef([resultAlias, alias]);
        }
      }
      throw new AggregateFunctionNotInSelectError(aggExpr.name);
    }
    case `func`: {
      const funcExpr = havingExpr;
      const transformedArgs = funcExpr.args.map(
        (arg) => replaceAggregatesByRefs(arg, selectClause)
      );
      return new Func(funcExpr.name, transformedArgs);
    }
    case `ref`:
      return havingExpr;
    case `val`:
      return havingExpr;
    default:
      throw new UnknownHavingExpressionTypeError(havingExpr.type);
  }
}
function evaluateWrappedAggregates(finalResults, aggregatedRow, wrappedAggExprs, fields) {
  for (const key of Object.keys(aggregatedRow)) {
    if (key.startsWith(fields.aggregatePrefix)) {
      finalResults[key] = aggregatedRow[key];
    }
  }
  for (let i = 0; i < fields.groupKeyRefs.length; i++) {
    finalResults[fields.groupKeyRefs[i]] = aggregatedRow[fields.groupValues[i]];
  }
  for (const [alias, evaluator] of Object.entries(wrappedAggExprs)) {
    finalResults[alias] = evaluator(
      getGroupEvaluationRow(aggregatedRow, fields, finalResults)
    );
  }
  for (const key of Object.keys(finalResults)) {
    if (key.startsWith(fields.aggregatePrefix) || fields.groupKeyRefs.includes(key)) {
      delete finalResults[key];
    }
  }
}
function containsAggregate(expr) {
  if (isConditionalSelect(expr)) {
    const branchHasAggregate = expr.branches.some(
      (branch) => containsAggregate(branch.condition) || containsAggregate(branch.value)
    );
    return branchHasAggregate || expr.defaultValue !== void 0 && containsAggregate(expr.defaultValue);
  }
  if (isNestedSelectObject(expr)) {
    return Object.values(expr).some(
      (value) => containsAggregate(value)
    );
  }
  if (!isExpressionLike(expr)) {
    return false;
  }
  if (expr.type === `agg`) {
    return true;
  }
  if (expr.type === `func` && `args` in expr) {
    return expr.args.some(
      (arg) => containsAggregate(arg)
    );
  }
  return false;
}
function extractAndReplaceAggregates(expr, counter, aggregatePrefix) {
  if (expr.type === `includesSubquery`) {
    return { transformed: expr, extracted: {} };
  }
  if (expr.type === `agg`) {
    const alias = `${aggregatePrefix}${counter.value++}`;
    return {
      transformed: new PropRef([`$selected`, alias]),
      extracted: { [alias]: expr }
    };
  }
  if (expr.type === `func`) {
    const allExtracted = {};
    const newArgs = expr.args.map((arg) => {
      const result = extractAndReplaceAggregates(arg, counter, aggregatePrefix);
      Object.assign(allExtracted, result.extracted);
      return result.transformed;
    });
    return {
      transformed: new Func(expr.name, newArgs),
      extracted: allExtracted
    };
  }
  if (isConditionalSelect(expr)) {
    const allExtracted = {};
    const branches = expr.branches.map((branch) => {
      const condition = extractAndReplaceAggregates(
        branch.condition,
        counter,
        aggregatePrefix
      );
      const value = extractAndReplaceAggregates(
        branch.value,
        counter,
        aggregatePrefix
      );
      Object.assign(allExtracted, condition.extracted, value.extracted);
      return {
        condition: condition.transformed,
        value: value.transformed
      };
    });
    const defaultValue = expr.defaultValue === void 0 ? void 0 : extractAndReplaceAggregates(
      expr.defaultValue,
      counter,
      aggregatePrefix
    );
    if (defaultValue) {
      Object.assign(allExtracted, defaultValue.extracted);
    }
    return {
      transformed: new ConditionalSelect(branches, defaultValue?.transformed),
      extracted: allExtracted
    };
  }
  if (isNestedSelectObject(expr)) {
    const allExtracted = {};
    const transformed = {};
    for (const [key, value] of Object.entries(expr)) {
      const result = extractAndReplaceAggregates(
        value,
        counter,
        aggregatePrefix
      );
      Object.assign(allExtracted, result.extracted);
      transformed[key] = result.transformed;
    }
    return { transformed, extracted: allExtracted };
  }
  return { transformed: expr, extracted: {} };
}
function replaceGroupByRefsInSelectValue(value, groupByClause, groupKeyRefs) {
  if (isConditionalSelect(value)) {
    return new ConditionalSelect(
      value.branches.map((branch) => ({
        condition: replaceGroupByRefsInExpression(
          branch.condition,
          groupByClause,
          groupKeyRefs
        ),
        value: replaceGroupByRefsInSelectValue(
          branch.value,
          groupByClause,
          groupKeyRefs
        )
      })),
      value.defaultValue === void 0 ? void 0 : replaceGroupByRefsInSelectValue(
        value.defaultValue,
        groupByClause,
        groupKeyRefs
      )
    );
  }
  if (isNestedSelectObject(value)) {
    const transformed = {};
    for (const [key, entry] of Object.entries(value)) {
      transformed[key] = replaceGroupByRefsInSelectValue(
        entry,
        groupByClause,
        groupKeyRefs
      );
    }
    return transformed;
  }
  if (!isExpressionLike(value)) {
    return value;
  }
  if (value.type === `includesSubquery` || value.type === `agg`) {
    return value;
  }
  return replaceGroupByRefsInExpression(value, groupByClause, groupKeyRefs);
}
function replaceGroupByRefsInExpression(expr, groupByClause, groupKeyRefs) {
  if (expr.type === `ref`) {
    const groupIndex = groupByClause.findIndex(
      (groupExpr) => expressionsEqual(expr, groupExpr)
    );
    return groupIndex === -1 ? expr : new PropRef([`$selected`, groupKeyRefs[groupIndex]]);
  }
  if (expr.type === `func`) {
    return new Func(
      expr.name,
      expr.args.map(
        (arg) => replaceGroupByRefsInExpression(arg, groupByClause, groupKeyRefs)
      )
    );
  }
  return expr;
}
function compileGroupedSelectValue(value) {
  if (isConditionalSelect(value)) {
    return compileGroupedConditionalSelect(value);
  }
  if (value.type === `includesSubquery`) {
    return () => null;
  }
  if (isNestedSelectObject(value)) {
    return compileGroupedSelectObject(value);
  }
  if (!isExpressionLike(value)) {
    return () => value;
  }
  return compileExpression(value);
}
function compileGroupedSelectObject(obj) {
  const entries = Object.entries(obj).map(([key, value]) => {
    if (key.startsWith(`__SPREAD_SENTINEL__`)) {
      const rest = key.slice(`__SPREAD_SENTINEL__`.length);
      const splitIndex = rest.lastIndexOf(`__`);
      const pathStr = splitIndex >= 0 ? rest.slice(0, splitIndex) : rest;
      const isRefExpr = typeof value === `object` && `type` in value && value.type === `ref`;
      const expression = isRefExpr ? value : new PropRef(pathStr.split(`.`));
      return {
        key,
        spread: true,
        value: compileExpression(expression)
      };
    }
    return {
      key,
      spread: false,
      value: compileGroupedSelectValue(value)
    };
  });
  return (row) => {
    const result = {};
    for (const entry of entries) {
      const value = entry.value(row);
      if (entry.spread) {
        if (value && typeof value === `object`) {
          Object.assign(result, value);
        }
      } else {
        result[entry.key] = value;
      }
    }
    return result;
  };
}
function compileGroupedConditionalSelect(conditional) {
  const branches = conditional.branches.map((branch) => ({
    condition: compileExpression(branch.condition),
    value: compileGroupedSelectValue(branch.value)
  }));
  const defaultValue = conditional.defaultValue === void 0 ? void 0 : compileGroupedSelectValue(conditional.defaultValue);
  return (row) => {
    for (const branch of branches) {
      if (isCaseWhenConditionTrue(branch.condition(row))) {
        return branch.value(row);
      }
    }
    return defaultValue !== void 0 ? defaultValue(row) : null;
  };
}
function isNestedSelectObject(value) {
  return value != null && typeof value === `object` && !Array.isArray(value) && !value.__refProxy && !isExpressionLike(value);
}
function isConditionalSelect(value) {
  return value instanceof ConditionalSelect || value != null && typeof value === `object` && value.type === `conditionalSelect`;
}
function aggregatesEqual(agg1, agg2) {
  return agg1.name === agg2.name && agg1.args.length === agg2.args.length && agg1.args.every((arg, i) => expressionsEqual(arg, agg2.args[i]));
}
export {
  containsAggregate,
  processGroupBy,
  replaceAggregatesByRefs
};
//# sourceMappingURL=group-by.js.map