@tanstack/db/dist/esm/query/compiler/joins.js

A reactive client store for building super fast apps on sync

import { map, tap, serializeValue, join, filter } from "@tanstack/db-ivm";
import { InvalidJoinConditionSourceMismatchError, JoinCollectionNotFoundError, UnsupportedJoinTypeError, InvalidJoinConditionSameSourceError, InvalidJoinConditionLeftSourceError, InvalidJoinConditionRightSourceError, InvalidJoinCondition, UnsupportedJoinSourceTypeError, CollectionInputNotFoundError } from "../../errors.js";
import { normalizeValue } from "../../utils/comparison.js";
import { ensureIndexForField } from "../../indexes/auto-index.js";
import { compileExpression } from "./evaluators.js";
import { getLazyLoadTargets } from "./lazy-targets.js";
import { crossJoinParentRoutes } from "./parent-routes.js";
import { attachRouteMetadataToResult, getRoutedScalarMetadata, INCLUDES_PUBLIC_KEY } from "./route-metadata.js";
let nextLazyDemandPlanId = 0;
function parameterizeJoinInputByParentRoutes(input, parentKeyStream) {
  return crossJoinParentRoutes(
    input,
    parentKeyStream,
    (rowKey, row, correlationKey, parentContext) => {
      return [
        serializeValue([rowKey, correlationKey, parentContext]),
        {
          ...row,
          __correlationKey: correlationKey,
          __parentContext: parentContext
        }
      ];
    }
  );
}
function wrapJoinedInputRow(alias, row) {
  const scalar = getRoutedScalarMetadata(row);
  if (scalar) {
    const namespaced2 = {
      [alias]: scalar.value,
      __correlationKey: scalar.correlationKey,
      __parentContext: scalar.parentContext,
      [INCLUDES_PUBLIC_KEY]: scalar.publicKey
    };
    if (scalar.parentContext != null && typeof scalar.parentContext === `object`) {
      Object.assign(namespaced2, scalar.parentContext);
    }
    return namespaced2;
  }
  if (row == null || typeof row !== `object`) {
    return { [alias]: row };
  }
  const { __parentContext, ...cleanRow } = row;
  const namespaced = { [alias]: cleanRow };
  if (__parentContext != null) {
    Object.assign(namespaced, __parentContext);
    namespaced.__parentContext = __parentContext;
  }
  return namespaced;
}
function getRouteJoinKey(row, source, value) {
  return serializeValue([
    row[source]?.__correlationKey ?? row.__correlationKey,
    row.__parentContext ?? row[source]?.__parentContext ?? null,
    value
  ]);
}
function registerLazyDemandPlan(callbacks, target, initialKeys = /* @__PURE__ */ new Set()) {
  const plan = {
    id: `lazy-demand-${++nextLazyDemandPlanId}`,
    path: target.path,
    collectionId: target.collection.id,
    initialKeys: new Set(initialKeys)
  };
  const state = callbacks[target.sourceId] ??= {};
  (state.plans ??= []).push(plan);
  return plan;
}
function processJoins(pipeline, joinClauses, sources, mainCollectionId, mainSource, allInputs, cache, queryMapping, collections, subscriptions, callbacks, lazySources, optimizableOrderByCollections, setWindowFn, rawQuery, onCompileSubquery, aliasToCollectionId, aliasRemapping, sourceWhereClauses, mainSourceIsParentFiltered, parentKeyStream) {
  let resultPipeline = pipeline;
  for (const joinClause of joinClauses) {
    resultPipeline = processJoin(
      resultPipeline,
      joinClause,
      sources,
      mainCollectionId,
      mainSource,
      allInputs,
      cache,
      queryMapping,
      collections,
      subscriptions,
      callbacks,
      lazySources,
      optimizableOrderByCollections,
      setWindowFn,
      rawQuery,
      onCompileSubquery,
      aliasToCollectionId,
      aliasRemapping,
      sourceWhereClauses,
      mainSourceIsParentFiltered,
      parentKeyStream
    );
  }
  return resultPipeline;
}
function processJoin(pipeline, joinClause, sources, mainCollectionId, mainSource, allInputs, cache, queryMapping, collections, subscriptions, callbacks, lazySources, optimizableOrderByCollections, setWindowFn, rawQuery, onCompileSubquery, aliasToCollectionId, aliasRemapping, sourceWhereClauses, mainSourceIsParentFiltered, parentKeyStream) {
  const isCollectionRef = joinClause.from.type === `collectionRef`;
  const joinedSource = joinClause.from.alias;
  const availableSources = [...Object.keys(sources), joinedSource];
  const { mainExpr, joinedExpr } = analyzeJoinExpressions(
    joinClause.left,
    joinClause.right,
    availableSources,
    joinedSource,
    rawQuery.from.type === `unionAll`
  );
  const joinedExpressionAliases = getSourceAliasesFromExpression(joinedExpr);
  const joinedExpressionUsesParent = [...joinedExpressionAliases].some(
    (alias) => alias !== joinedSource && !sources[alias]
  );
  const routeJoinedSource = parentKeyStream !== void 0 && (joinClause.from.type === `queryRef` || joinedExpressionUsesParent);
  const {
    alias: processedJoinedSource,
    input: joinedInput,
    collectionId: joinedCollectionId
  } = processJoinSource(
    joinClause.from,
    allInputs,
    collections,
    subscriptions,
    callbacks,
    lazySources,
    optimizableOrderByCollections,
    setWindowFn,
    cache,
    queryMapping,
    onCompileSubquery,
    aliasToCollectionId,
    aliasRemapping,
    sourceWhereClauses,
    routeJoinedSource ? parentKeyStream : void 0
  );
  if (processedJoinedSource !== joinedSource) {
    throw new InvalidJoinConditionSourceMismatchError();
  }
  sources[joinedSource] = joinedInput;
  if (isCollectionRef) {
    aliasToCollectionId[joinedSource] = joinedCollectionId;
  }
  const mainCollection = collections[mainCollectionId];
  const joinedCollection = collections[joinedCollectionId];
  if (!mainCollection) {
    throw new JoinCollectionNotFoundError(mainCollectionId);
  }
  if (!joinedCollection) {
    throw new JoinCollectionNotFoundError(joinedCollectionId);
  }
  const sourceActivity = getActiveAndLazySources(
    joinClause.type,
    mainCollection,
    joinedCollection,
    mainSourceIsParentFiltered
  );
  const activeSource = routeJoinedSource ? void 0 : sourceActivity.activeSource;
  const lazySource = sourceActivity.lazySource;
  const compiledMainExpr = compileExpression(mainExpr);
  const compiledJoinedExpr = compileExpression(joinedExpr);
  let mainPipeline = pipeline.pipe(
    map(([currentKey, namespacedRow]) => {
      const value = normalizeValue(compiledMainExpr(namespacedRow));
      const mainKey = routeJoinedSource ? getRouteJoinKey(namespacedRow, mainSource, value) : value;
      return [mainKey, [currentKey, namespacedRow]];
    })
  );
  let joinedPipeline = joinedInput.pipe(
    map(([currentKey, row]) => {
      const namespacedRow = wrapJoinedInputRow(joinedSource, row);
      const value = normalizeValue(compiledJoinedExpr(namespacedRow));
      const joinedKey = routeJoinedSource ? getRouteJoinKey(namespacedRow, joinedSource, value) : value;
      return [joinedKey, [currentKey, namespacedRow]];
    })
  );
  if (![`inner`, `left`, `right`, `full`].includes(joinClause.type)) {
    throw new UnsupportedJoinTypeError(joinClause.type);
  }
  if (activeSource) {
    const lazyFrom = activeSource === `main` ? joinClause.from : rawQuery.from;
    const limitedSubquery = lazyFrom.type === `queryRef` && (lazyFrom.query.limit || lazyFrom.query.offset);
    const resultUnionLazySide = lazyFrom.type === `unionAll`;
    const lazySourceJoinExpr = activeSource === `main` ? joinedExpr : mainExpr;
    const lazyAlias = activeSource === `main` ? joinedSource : mainSource;
    const lazyTargets = resultUnionLazySide ? [] : getLazyLoadTargets(
      rawQuery,
      lazyFrom,
      lazyAlias,
      lazySourceJoinExpr,
      lazySource,
      aliasRemapping
    );
    if (!limitedSubquery && lazyTargets.length > 0) {
      for (const target of lazyTargets) {
        lazySources.add(target.sourceId);
      }
      const demandPlans = lazyTargets.map(
        (target) => registerLazyDemandPlan(callbacks, target)
      );
      const demandWeights = /* @__PURE__ */ new Map();
      const activePipeline = activeSource === `main` ? mainPipeline : joinedPipeline;
      for (const target of lazyTargets) {
        const fieldName = target.path[0];
        if (fieldName) {
          ensureIndexForField(fieldName, target.path, target.collection);
        }
      }
      const activePipelineWithLoading = activePipeline.pipe(
        tap((data) => {
          for (const [[joinKey], weight] of data.getInner()) {
            if (joinKey == null) continue;
            const encoded = serializeValue(joinKey);
            const previous = demandWeights.get(encoded);
            const nextWeight = (previous?.weight ?? 0) + weight;
            if (nextWeight === 0) {
              demandWeights.delete(encoded);
            } else {
              demandWeights.set(encoded, { key: joinKey, weight: nextWeight });
            }
          }
          const keys = new Set(
            [...demandWeights.values()].filter(({ weight }) => weight > 0).map(({ key }) => key)
          );
          for (let index = 0; index < lazyTargets.length; index++) {
            const target = lazyTargets[index];
            callbacks[target.sourceId]?.setDemand?.(demandPlans[index], keys);
          }
        })
      );
      if (activeSource === `main`) {
        mainPipeline = activePipelineWithLoading;
      } else {
        joinedPipeline = activePipelineWithLoading;
      }
    }
  }
  return mainPipeline.pipe(
    join(joinedPipeline, joinClause.type),
    processJoinResults(joinClause.type)
  );
}
function analyzeJoinExpressions(left, right, allAvailableSourceAliases, joinedSource, allowResultFields = false) {
  const availableSources = allAvailableSourceAliases.filter(
    (alias) => alias !== joinedSource
  );
  const leftSourceAliases = getSourceAliasesFromExpression(left);
  const rightSourceAliases = getSourceAliasesFromExpression(right);
  const leftReferencesJoined = leftSourceAliases.has(joinedSource);
  const rightReferencesJoined = rightSourceAliases.has(joinedSource);
  const leftAvailableAliases = [...leftSourceAliases].filter(
    (alias) => availableSources.includes(alias) || allowResultFields && alias !== joinedSource
  );
  const rightAvailableAliases = [...rightSourceAliases].filter(
    (alias) => availableSources.includes(alias) || allowResultFields && alias !== joinedSource
  );
  if (leftAvailableAliases.length > 0 && !leftReferencesJoined && rightReferencesJoined && rightAvailableAliases.length === 0) {
    return { mainExpr: left, joinedExpr: right };
  }
  if (leftReferencesJoined && leftAvailableAliases.length === 0 && rightAvailableAliases.length > 0 && !rightReferencesJoined) {
    return { mainExpr: right, joinedExpr: left };
  }
  if (leftSourceAliases.size === 0 || rightSourceAliases.size === 0) {
    throw new InvalidJoinConditionSourceMismatchError();
  }
  if (leftSourceAliases.size === 1 && rightSourceAliases.size === 1 && [...leftSourceAliases][0] === [...rightSourceAliases][0]) {
    throw new InvalidJoinConditionSameSourceError([...leftSourceAliases][0]);
  }
  if (leftAvailableAliases.length === 0) {
    throw new InvalidJoinConditionLeftSourceError([...leftSourceAliases][0]);
  }
  if (!rightReferencesJoined) {
    throw new InvalidJoinConditionRightSourceError(joinedSource);
  }
  throw new InvalidJoinCondition();
}
function getSourceAliasesFromExpression(expr) {
  switch (expr.type) {
    case `ref`:
      return new Set(expr.path[0] ? [expr.path[0]] : []);
    case `func`: {
      const sourceAliases = /* @__PURE__ */ new Set();
      for (const arg of expr.args) {
        for (const alias of getSourceAliasesFromExpression(arg)) {
          sourceAliases.add(alias);
        }
      }
      return sourceAliases;
    }
    default:
      return /* @__PURE__ */ new Set();
  }
}
function processJoinSource(from, allInputs, collections, subscriptions, callbacks, lazySources, optimizableOrderByCollections, setWindowFn, cache, queryMapping, onCompileSubquery, aliasToCollectionId, aliasRemapping, sourceWhereClauses, parentKeyStream) {
  switch (from.type) {
    case `collectionRef`: {
      const input = allInputs[from.sourceId] ?? allInputs[from.alias];
      if (!input) {
        throw new CollectionInputNotFoundError(
          from.alias,
          from.collection.id,
          Object.keys(allInputs)
        );
      }
      aliasToCollectionId[from.alias] = from.collection.id;
      return {
        alias: from.alias,
        input: parentKeyStream ? parameterizeJoinInputByParentRoutes(input, parentKeyStream) : input,
        collectionId: from.collection.id
      };
    }
    case `queryRef`: {
      const originalQuery = queryMapping.get(from.query) || from.query;
      const subQueryResult = onCompileSubquery(
        originalQuery,
        allInputs,
        collections,
        subscriptions,
        callbacks,
        lazySources,
        optimizableOrderByCollections,
        setWindowFn,
        cache,
        queryMapping,
        parentKeyStream
      );
      Object.assign(aliasToCollectionId, subQueryResult.aliasToCollectionId);
      Object.assign(aliasRemapping, subQueryResult.aliasRemapping);
      const isUserDefinedSubquery = queryMapping.has(from.query);
      const fromInnerAlias = getFirstFromAlias(from.query);
      const isOptimizerCreated = !isUserDefinedSubquery && fromInnerAlias !== void 0 && from.alias === fromInnerAlias;
      if (!isOptimizerCreated) {
        for (const [alias, whereClause] of subQueryResult.sourceWhereClauses) {
          sourceWhereClauses.set(alias, whereClause);
        }
      }
      const innerAlias = Object.keys(subQueryResult.aliasToCollectionId).find(
        (alias) => subQueryResult.aliasToCollectionId[alias] === subQueryResult.collectionId
      );
      if (innerAlias && innerAlias !== from.alias) {
        aliasRemapping[from.alias] = innerAlias;
      }
      const subQueryInput = subQueryResult.pipeline;
      const extractedInput = subQueryInput.pipe(
        map((data) => {
          const [
            key,
            [value, _orderByIndex, correlationKey, parentContext, publicKey]
          ] = data;
          if (!parentKeyStream) {
            return [key, value];
          }
          return [
            key,
            attachRouteMetadataToResult(
              value,
              correlationKey,
              parentContext,
              publicKey
            )
          ];
        })
      );
      return {
        alias: from.alias,
        input: extractedInput,
        collectionId: subQueryResult.collectionId
      };
    }
    default:
      throw new UnsupportedJoinSourceTypeError(from.type);
  }
}
function getFirstFromAlias(query) {
  if (query.from.type === `unionFrom`) {
    return query.from.sources[0]?.alias;
  }
  if (query.from.type === `unionAll`) {
    return void 0;
  }
  return query.from.alias;
}
function processJoinResults(joinType) {
  return function(pipeline) {
    return pipeline.pipe(
      // Process the join result and handle nulls
      filter((result) => {
        const [_key, [main, joined]] = result;
        const mainNamespacedRow = main?.[1];
        const joinedNamespacedRow = joined?.[1];
        if (joinType === `inner`) {
          return !!(mainNamespacedRow && joinedNamespacedRow);
        }
        if (joinType === `left`) {
          return !!mainNamespacedRow;
        }
        if (joinType === `right`) {
          return !!joinedNamespacedRow;
        }
        return true;
      }),
      map((result) => {
        const [_key, [main, joined]] = result;
        const mainKey = main?.[0];
        const mainNamespacedRow = main?.[1];
        const joinedKey = joined?.[0];
        const joinedNamespacedRow = joined?.[1];
        const mergedNamespacedRow = {};
        if (mainNamespacedRow) {
          Object.assign(mergedNamespacedRow, mainNamespacedRow);
        }
        if (joinedNamespacedRow) {
          Object.assign(mergedNamespacedRow, joinedNamespacedRow);
        }
        const resultKey = `[${mainKey},${joinedKey}]`;
        return [resultKey, mergedNamespacedRow];
      })
    );
  };
}
function getActiveAndLazySources(joinType, leftCollection, rightCollection, mainSourceIsParentFiltered) {
  switch (joinType) {
    case `left`:
      return { activeSource: `main`, lazySource: rightCollection };
    case `right`:
      return { activeSource: `joined`, lazySource: leftCollection };
    case `inner`:
      if (mainSourceIsParentFiltered) {
        return { activeSource: `main`, lazySource: rightCollection };
      }
      return leftCollection.size < rightCollection.size ? { activeSource: `main`, lazySource: rightCollection } : { activeSource: `joined`, lazySource: leftCollection };
    default:
      return { activeSource: void 0, lazySource: void 0 };
  }
}
export {
  processJoins,
  registerLazyDemandPlan
};
//# sourceMappingURL=joins.js.map