node-redis/packages/search/lib/commands/AGGREGATE.ts

import { RedisCommandArgument, RedisCommandArguments } from '@node-redis/client/dist/lib/commands';
import { pushVerdictArgument, transformTuplesReply } from '@node-redis/client/dist/lib/commands/generic-transformers';
import { Params, PropertyName, pushArgumentsWithLength, pushParamsArgs, pushSortByArguments, SortByProperty } from '.';

export enum AggregateSteps {
    GROUPBY = 'GROUPBY',
    SORTBY = 'SORTBY',
    APPLY = 'APPLY',
    LIMIT = 'LIMIT',
    FILTER = 'FILTER'
}

interface AggregateStep<T extends AggregateSteps> {
    type: T;
}

export enum AggregateGroupByReducers {
    COUNT = 'COUNT',
    COUNT_DISTINCT = 'COUNT_DISTINCT',
    COUNT_DISTINCTISH = 'COUNT_DISTINCTISH',
    SUM = 'SUM',
    MIN = 'MIN',
    MAX = 'MAX',
    AVG = 'AVG',
    STDDEV = 'STDDEV',
    QUANTILE = 'QUANTILE',
    TOLIST = 'TOLIST',
    TO_LIST = 'TOLIST',
    FIRST_VALUE = 'FIRST_VALUE',
    RANDOM_SAMPLE = 'RANDOM_SAMPLE'
}

interface GroupByReducer<T extends AggregateGroupByReducers> {
    type: T;
    AS?: string;
}

type CountReducer = GroupByReducer<AggregateGroupByReducers.COUNT>;

interface CountDistinctReducer extends GroupByReducer<AggregateGroupByReducers.COUNT_DISTINCT> {
    property: PropertyName;
}

interface CountDistinctishReducer extends GroupByReducer<AggregateGroupByReducers.COUNT_DISTINCTISH> {
    property: PropertyName;
}

interface SumReducer extends GroupByReducer<AggregateGroupByReducers.SUM> {
    property: PropertyName;
}

interface MinReducer extends GroupByReducer<AggregateGroupByReducers.MIN> {
    property: PropertyName;
}

interface MaxReducer extends GroupByReducer<AggregateGroupByReducers.MAX> {
    property: PropertyName;
}

interface AvgReducer extends GroupByReducer<AggregateGroupByReducers.AVG> {
    property: PropertyName;
}

interface StdDevReducer extends GroupByReducer<AggregateGroupByReducers.STDDEV> {
    property: PropertyName;
}

interface QuantileReducer extends GroupByReducer<AggregateGroupByReducers.QUANTILE> {
    property: PropertyName;
    quantile: number;
}

interface ToListReducer extends GroupByReducer<AggregateGroupByReducers.TOLIST> {
    property: PropertyName;
}

interface FirstValueReducer extends GroupByReducer<AggregateGroupByReducers.FIRST_VALUE> {
    property: PropertyName;
    BY?: PropertyName | {
        property: PropertyName;
        direction?: 'ASC' | 'DESC';
    };
}

interface RandomSampleReducer extends GroupByReducer<AggregateGroupByReducers.RANDOM_SAMPLE> {
    property: PropertyName;
    sampleSize: number;
}

type GroupByReducers = CountReducer | CountDistinctReducer | CountDistinctishReducer | SumReducer | MinReducer | MaxReducer | AvgReducer | StdDevReducer | QuantileReducer | ToListReducer | FirstValueReducer | RandomSampleReducer;

interface GroupByStep extends AggregateStep<AggregateSteps.GROUPBY> {
    properties?: PropertyName | Array<PropertyName>;
    REDUCE: GroupByReducers | Array<GroupByReducers>;
}

interface SortStep extends AggregateStep<AggregateSteps.SORTBY> {
    BY: SortByProperty | Array<SortByProperty>;
    MAX?: number;
}

interface ApplyStep extends AggregateStep<AggregateSteps.APPLY> {
    expression: string;
    AS: string;
}

interface LimitStep extends AggregateStep<AggregateSteps.LIMIT> {
    from: number;
    size: number;
}

interface FilterStep extends AggregateStep<AggregateSteps.FILTER> {
    expression: string;
}

type LoadField = PropertyName | {
    identifier: PropertyName;
    AS?: string;
}

export interface AggregateOptions {
    VERBATIM?: true;
    LOAD?: LoadField | Array<LoadField>;
    STEPS?: Array<GroupByStep | SortStep | ApplyStep | LimitStep | FilterStep>;
    PARAMS?: Params;
    DIALECT?: number;
}

export function transformArguments(
    index: string,
    query: string,
    options?: AggregateOptions
): RedisCommandArguments {
    return pushAggregatehOptions(
        ['FT.AGGREGATE', index, query],
        options
    );
}

export function pushAggregatehOptions(
    args: RedisCommandArguments,
    options?: AggregateOptions
): RedisCommandArguments {
    if (options?.VERBATIM) {
        args.push('VERBATIM');
    }

    if (options?.LOAD) {
        args.push('LOAD');
        pushArgumentsWithLength(args, () => {
            if (Array.isArray(options.LOAD)) {
                for (const load of options.LOAD) {
                    pushLoadField(args, load);
                }
            } else {
                pushLoadField(args, options.LOAD!);
            }
        });
    }

    if (options?.STEPS) {
        for (const step of options.STEPS) {
            switch (step.type) {
                case AggregateSteps.GROUPBY:
                    args.push('GROUPBY');
                    if (!step.properties) {
                        args.push('0');
                    } else {
                        pushVerdictArgument(args, step.properties);
                    }

                    if (Array.isArray(step.REDUCE)) {
                        for (const reducer of step.REDUCE) {
                            pushGroupByReducer(args, reducer);
                        }
                    } else {
                        pushGroupByReducer(args, step.REDUCE);
                    }

                    break;

                case AggregateSteps.SORTBY:
                    pushSortByArguments(args, 'SORTBY', step.BY);

                    if (step.MAX) {
                        args.push('MAX', step.MAX.toString());
                    }

                    break;

                case AggregateSteps.APPLY:
                    args.push('APPLY', step.expression, 'AS', step.AS);
                    break;

                case AggregateSteps.LIMIT:
                    args.push('LIMIT', step.from.toString(), step.size.toString());
                    break;

                case AggregateSteps.FILTER:
                    args.push('FILTER', step.expression);
                    break;
            }
        }
    }

    pushParamsArgs(args, options?.PARAMS);

    if (options?.DIALECT) {
        args.push('DIALECT', options.DIALECT.toString());
    }

    return args;
}

function pushLoadField(args: RedisCommandArguments, toLoad: LoadField): void {
    if (typeof toLoad === 'string') {
        args.push(toLoad);
    } else {
        args.push(toLoad.identifier);

        if (toLoad.AS) {
            args.push('AS', toLoad.AS);
        }
    }
}

function pushGroupByReducer(args: RedisCommandArguments, reducer: GroupByReducers): void {
    args.push('REDUCE', reducer.type);

    switch (reducer.type) {
        case AggregateGroupByReducers.COUNT:
            args.push('0');
            break;

        case AggregateGroupByReducers.COUNT_DISTINCT:
        case AggregateGroupByReducers.COUNT_DISTINCTISH:
        case AggregateGroupByReducers.SUM:
        case AggregateGroupByReducers.MIN:
        case AggregateGroupByReducers.MAX:
        case AggregateGroupByReducers.AVG:
        case AggregateGroupByReducers.STDDEV:
        case AggregateGroupByReducers.TOLIST:
            args.push('1', reducer.property);
            break;

        case AggregateGroupByReducers.QUANTILE:
            args.push('2', reducer.property, reducer.quantile.toString());
            break;

        case AggregateGroupByReducers.FIRST_VALUE: {
            pushArgumentsWithLength(args, () => {
                args.push(reducer.property);

                if (reducer.BY) {
                    args.push('BY');
                    if (typeof reducer.BY === 'string') {
                        args.push(reducer.BY);
                    } else {
                        args.push(reducer.BY.property);

                        if (reducer.BY.direction) {
                            args.push(reducer.BY.direction);
                        }
                    }
                }
            });
            break;
        }

        case AggregateGroupByReducers.RANDOM_SAMPLE:
            args.push('2', reducer.property, reducer.sampleSize.toString());
            break;
    }

    if (reducer.AS) {
        args.push('AS', reducer.AS);
    }
}

export type AggregateRawReply = [
    total: number,
    ...results: Array<Array<RedisCommandArgument>>
];

export interface AggregateReply {
    total: number;
    results: Array<Record<string, RedisCommandArgument>>;
}

export function transformReply(rawReply: AggregateRawReply): AggregateReply {
    const results: Array<Record<string, RedisCommandArgument>> = [];
    for (let i = 1; i < rawReply.length; i++) {
        results.push(
            transformTuplesReply(rawReply[i] as Array<RedisCommandArgument>)
        );
    }

    return {
        total: rawReply[0],
        results
    };
}