mingo/esm/operators/pipeline/setWindowFields.js

MongoDB query language for in-memory objects

import { ComputeOptions, OpType } from "../../core/_internal";
import { concat, Lazy } from "../../lazy";
import { assert, isNumber, isOperator, isString } from "../../util";
import { $function } from "../expression/custom/function";
import { $dateAdd } from "../expression/date/dateAdd";
import { $addFields } from "./addFields";
import { $group } from "./group";
import { $sort } from "./sort";
const SORT_REQUIRED_OPS = /* @__PURE__ */ new Set([
  "$denseRank",
  "$documentNumber",
  "$first",
  "$last",
  "$linearFill",
  "$rank",
  "$shift"
]);
const WINDOW_UNBOUNDED_OPS = /* @__PURE__ */ new Set([
  "$denseRank",
  "$expMovingAvg",
  "$linearFill",
  "$locf",
  "$rank",
  "$shift"
]);
const isUnbounded = (window) => {
  const boundary = window?.documents || window?.range;
  return !boundary || boundary[0] === "unbounded" && boundary[1] === "unbounded";
};
const $setWindowFields = (collection, expr, options) => {
  options = ComputeOptions.init(options);
  options.context.addExpressionOps({ $function });
  for (const outputExpr of Object.values(expr.output)) {
    const keys = Object.keys(outputExpr);
    const op = keys.find(isOperator);
    const context = options.context;
    assert(
      !!context.getOperator(OpType.WINDOW, op) || !!context.getOperator(OpType.ACCUMULATOR, op),
      `'${op}' is not a valid window operator`
    );
    assert(
      keys.length > 0 && keys.length <= 2 && (keys.length == 1 || keys.includes("window")),
      "'output' option should have a single window operator."
    );
    if (outputExpr?.window) {
      const { documents, range } = outputExpr.window;
      assert(
        (+!!documents ^ +!!range) === 1,
        "'window' option supports only one of 'documents' or 'range'."
      );
    }
  }
  if (expr.sortBy) {
    collection = $sort(collection, expr.sortBy, options);
  }
  collection = $group(
    collection,
    {
      _id: expr.partitionBy,
      items: { $push: "$$CURRENT" }
    },
    options
  );
  return collection.transform(((partitions) => {
    const iterators = [];
    const outputConfig = [];
    for (const [field, outputExpr] of Object.entries(expr.output)) {
      const op = Object.keys(outputExpr).find(isOperator);
      const config = {
        operatorName: op,
        func: {
          left: options.context.getOperator(
            OpType.ACCUMULATOR,
            op
          ),
          right: options.context.getOperator(
            OpType.WINDOW,
            op
          )
        },
        args: outputExpr[op],
        field,
        window: outputExpr.window
      };
      assert(
        !!expr.sortBy || !(SORT_REQUIRED_OPS.has(op) || !config.window),
        `${SORT_REQUIRED_OPS.has(op) ? `'${op}'` : "bounded window operation"} requires a sortBy.`
      );
      assert(
        !config.window || !WINDOW_UNBOUNDED_OPS.has(op),
        `${op} does not accept a 'window' field.`
      );
      outputConfig.push(config);
    }
    partitions.forEach(((group) => {
      const items = group.items;
      let iterator = Lazy(items);
      const windowResultMap = {};
      for (const config of outputConfig) {
        const { func, args, field, window } = config;
        const makeResultFunc = (getItemsFn) => {
          let index = -1;
          return (obj) => {
            ++index;
            if (func.left) {
              return func.left(getItemsFn(obj, index), args, options);
            } else if (func.right) {
              return func.right(
                obj,
                getItemsFn(obj, index),
                {
                  parentExpr: expr,
                  inputExpr: args,
                  documentNumber: index + 1,
                  field
                },
                // must use raw options only since it operates over a collection.
                options
              );
            }
          };
        };
        if (window) {
          const { documents, range, unit } = window;
          const boundary = documents || range;
          if (!isUnbounded(window)) {
            const [begin, end] = boundary;
            const toBeginIndex = (currentIndex) => {
              if (begin == "current") return currentIndex;
              if (begin == "unbounded") return 0;
              return Math.max(begin + currentIndex, 0);
            };
            const toEndIndex = (currentIndex) => {
              if (end == "current") return currentIndex + 1;
              if (end == "unbounded") return items.length;
              return end + currentIndex + 1;
            };
            const getItems = (current, index) => {
              if (!!documents || boundary.every(isString)) {
                return items.slice(toBeginIndex(index), toEndIndex(index));
              }
              const sortKey = Object.keys(expr.sortBy)[0];
              let lower;
              let upper;
              if (unit) {
                const getTime = (amount) => {
                  return $dateAdd(
                    current,
                    {
                      startDate: new Date(current[sortKey]),
                      unit,
                      amount
                    },
                    options
                  ).getTime();
                };
                lower = isNumber(begin) ? getTime(begin) : -Infinity;
                upper = isNumber(end) ? getTime(end) : Infinity;
              } else {
                const currentValue = current[sortKey];
                lower = isNumber(begin) ? currentValue + begin : -Infinity;
                upper = isNumber(end) ? currentValue + end : Infinity;
              }
              let i = begin == "current" ? index : 0;
              const sliceEnd = end == "current" ? index + 1 : items.length;
              const array = new Array();
              while (i < sliceEnd) {
                const o = items[i++];
                const n = +o[sortKey];
                if (n >= lower && n <= upper) array.push(o);
              }
              return array;
            };
            windowResultMap[field] = makeResultFunc(getItems);
          }
        }
        if (!windowResultMap[field]) {
          windowResultMap[field] = makeResultFunc((_) => items);
        }
        iterator = $addFields(
          iterator,
          {
            [field]: {
              $function: {
                body: (obj) => windowResultMap[field](obj),
                args: ["$$CURRENT"]
              }
            }
          },
          options
        );
      }
      iterators.push(iterator);
    }));
    return concat(...iterators);
  }));
};
export {
  $setWindowFields
};