@uwdata/mosaic-core/src/preagg/preagg-columns.ts

Scalable and extensible linked data views.

import type { AggregateNode, ColumnRefNode, ExprNode, Query, SelectQuery } from '@uwdata/mosaic-sql';
import type { MosaicClient } from '../MosaicClient.js';
import { collectAggregates, FromClauseNode, isAggregateExpression, isColumnRef, isSelectQuery, isTableRef, rewrite, sql } from '@uwdata/mosaic-sql';
import { sufficientStatistics } from './sufficient-statistics.js';

export interface PreAggColumnsResult {
  dims: string[];
  groupby: Record<string, ExprNode>;
  preagg: Record<string, ExprNode>;
  output: Record<string, ExprNode>;
  having: ExprNode[];
  qualify: ExprNode[];
}

/**
 * Determine pre-aggregation columns for a given Mosaic client.
 * @param client The Mosaic client.
 * @returns An object with necessary column data to generate pre-aggregated
 *  columns, or null if the client can't be optimized or the client query
 *  contains an invalid or unsupported expression.
 */
export function preaggColumns(client: MosaicClient): PreAggColumnsResult | null {
  if (!client.filterStable) return null;
  const q = client.query();

  // bail if query is not analyzable
  if (!isSelectQuery(q)) return null;

  // bail if no base table
  const from = getBase(q, q => {
    const node = q._from[0];
    const ref = node instanceof FromClauseNode && node.expr;
    return isTableRef(ref) ? ref.name : ref;
  });
  if (typeof from !== 'string') return null;

  const aggrs = new Map<AggregateNode, ExprNode>();
  const preagg: Record<string, ExprNode> = {};
  const output: Record<string, ExprNode> = {};

  // groupby entries, these may or may not be selected
  const groupby: Record<string, ExprNode> = {};
  for (const expr of q._groupby) {
    // ignore integer index, as expr will be in select clause
    // we will harvest that expr as a dimension below
    if (expr.type !== "LITERAL") {
      const alias = isColumnRef(expr) ? expr.column : `${expr}`;
      groupby[alias] = expr;
    }
  }

  // all group by dimensions, including selected ones.
  const dims: Set<string> = new Set(Object.keys(groupby));

  // generate a scalar subquery for a global average
  const avg = (ref: ColumnRefNode) => {
    const name = ref.column;
    const expr = getBase(q, q => q._select.find(c => c.alias === name)?.expr);
    return sql`(SELECT avg(${expr ?? ref}) FROM "${from}")`;
  };

  try {
    // iterate over select clauses and analyze expressions
    for (const { alias, expr } of q._select) {
      const result = analyzeExpression(expr, aggrs, preagg, avg);
      if (result) {
        // rewrite original select clause to use preaggregates
        output[alias] = result;
      } else {
        // if no aggregates, expr is a groupby dimension
        dims.add(alias);
        preagg[alias] = expr;
      }
    }

    // bail if select clause has no aggregates
    if (!aggrs.size) return null;

    // analyze any having expressions
    const having = q._having
      .map(expr => analyzeExpression(expr, aggrs, preagg, avg) ?? expr);

    // analyze any qualify expressions
    const qualify = q._qualify
      .map(expr => analyzeExpression(expr, aggrs, preagg, avg) ?? expr);

    // bail if the query has no aggregates
    if (!aggrs.size) return null;

    return {
      dims: Array.from(dims),
      groupby,
      preagg,
      output,
      having,
      qualify
    };
  } catch {
    // bail if unsupported aggregate was encountered
    return null;
  }
}

/**
 * Analyzes an expression and returns a rewritten expression if preaggregation
 * optimization is supported. Returns undefined if the expressions does not
 * contain any aggregates. Throws if the expression contains an unsupported
 * or non-optimizable aggregate.
 */
function analyzeExpression(
  expr: ExprNode,
  aggrs: Map<AggregateNode, ExprNode>,
  preagg: Record<string, ExprNode>,
  avg: (ref: ColumnRefNode) => ExprNode
) {
  // bail if there is an aggregate we can't analyze
  // a value > 1 indicates an aggregate in verbatim text
  if (isAggregateExpression(expr!) > 1) throw new Error();

  const nodes = collectAggregates(expr!);
  if (nodes.length) {
    for (const node of nodes) {
      // bail if distinct aggregate
      if (node.isDistinct) throw new Error();

      // bail if aggregate function is unsupported
      // otherwise add output aggregate to rewrite map
      const agg = sufficientStatistics(node, preagg, avg);
      if (!agg) throw new Error();
      aggrs.set(node, agg);
    }

    // rewrite original expression to use preaggregates
    return rewrite(expr!, aggrs)!;
  }
}

/**
 * Identify a shared base (source) query and extract a value from it.
 * This method is used to find a shared base table name or extract
 * the original column name within a base table.
 * @param query The input query.
 * @param get A getter function to extract
 *  a value from a base query.
 * @returns the base query value, or
 *  `undefined` if there is no source table, or `NaN` if the
 *  query operates over multiple source tables.
 */
function getBase<T>(
  query: Query,
  get: (q: SelectQuery) => T
): T | number | undefined {
  const subq = query.subqueries;

  // select query
  if (isSelectQuery(query) && subq.length === 0) {
    return get(query);
  }

  // handle set operations / subqueries
  const base = getBase(subq[0], get);
  for (let i = 1; i < subq.length; ++i) {
    const value = getBase(subq[i], get);
    if (value === undefined) continue;
    if (value !== base) return NaN;
  }
  return base;
}