Skip to content

feat(search): add COLLECT reducer to FT.AGGREGATE #3310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nkaradzhov wants to merge 1 commit into
base: master
Choose a base branch
from
+245 −2
Open

feat(search): add COLLECT reducer to FT.AGGREGATE #3310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
190 changes: 190 additions & 0 deletions packages/search/lib/commands/AGGREGATE.spec.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,162 @@ describe('AGGREGATE', () => {
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'RANDOM_SAMPLE', '2', '@property', '1', 'DIALECT', DEFAULT_DIALECT]
);
});

describe('COLLECT', () => {
it('FIELDS *', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '*'
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '2', 'FIELDS', '*', 'DIALECT', DEFAULT_DIALECT]
);
});

it('FIELDS single field', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '@field'
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '3', 'FIELDS', '1', '@field', 'DIALECT', DEFAULT_DIALECT]
);
});

it('FIELDS array', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: ['@a', '@b']
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '4', 'FIELDS', '2', '@a', '@b', 'DIALECT', DEFAULT_DIALECT]
);
});

it('with DISTINCT', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
DISTINCT: true
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '3', 'FIELDS', '*', 'DISTINCT', 'DIALECT', DEFAULT_DIALECT]
);
});

it('with SORTBY string', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
SORTBY: '@field'
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '5', 'FIELDS', '*', 'SORTBY', '1', '@field', 'DIALECT', DEFAULT_DIALECT]
);
});

it('with SORTBY array and directions', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
SORTBY: [
{ BY: '@target', DIRECTION: 'DESC' },
{ BY: '@bestByDate', DIRECTION: 'ASC' }
]
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '8', 'FIELDS', '*', 'SORTBY', '4', '@target', 'DESC', '@bestByDate', 'ASC', 'DIALECT', DEFAULT_DIALECT]
);
});

it('with LIMIT', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
LIMIT: { from: 0, size: 5 }
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '0', 'REDUCE', 'COLLECT', '5', 'FIELDS', '*', 'LIMIT', '0', '5', 'DIALECT', DEFAULT_DIALECT]
);
});

it('movies-by-genre example (narg 9)', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
properties: '@genre',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
SORTBY: { BY: '@rating', DIRECTION: 'DESC' },
LIMIT: { from: 0, size: 5 },
AS: 'top_movies'
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '1', '@genre', 'REDUCE', 'COLLECT', '9', 'FIELDS', '*', 'SORTBY', '2', '@rating', 'DESC', 'LIMIT', '0', '5', 'AS', 'top_movies', 'DIALECT', DEFAULT_DIALECT]
);
});

it('relationships example (narg 12, DISTINCT)', () => {
assert.deepEqual(
parseArgs(AGGREGATE, 'index', '*', {
STEPS: [{
type: 'GROUPBY',
properties: '@relationshipName',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
DISTINCT: true,
SORTBY: [
{ BY: '@target', DIRECTION: 'DESC' },
{ BY: '@bestByDate', DIRECTION: 'ASC' }
],
LIMIT: { from: 0, size: 3 },
AS: 'all_distinct_docs'
}
}]
}),
['FT.AGGREGATE', 'index', '*', 'GROUPBY', '1', '@relationshipName', 'REDUCE', 'COLLECT', '12', 'FIELDS', '*', 'DISTINCT', 'SORTBY', '4', '@target', 'DESC', '@bestByDate', 'ASC', 'LIMIT', '0', '3', 'AS', 'all_distinct_docs', 'DIALECT', DEFAULT_DIALECT]
);
});
});
});

describe('SORTBY', () => {
Expand Down Expand Up @@ -602,4 +758,38 @@ describe('AGGREGATE', () => {
}
}
});

testUtils.testWithClient('client.ft.aggregate with COLLECT reducer', async client => {
await client.ft.create('index', {
genre: 'TAG',
rating: 'NUMERIC'
});
await client.hSet('movie:1', { genre: 'action', rating: '8' });
await client.hSet('movie:2', { genre: 'action', rating: '9' });
await client.hSet('movie:3', { genre: 'drama', rating: '7' });

const reply = await client.ft.aggregate('index', '*', {
LOAD: '*',
STEPS: [{
type: 'GROUPBY',
properties: '@genre',
REDUCE: {
type: 'COLLECT',
FIELDS: '*',
SORTBY: { BY: '@rating', DIRECTION: 'DESC' },
LIMIT: { from: 0, size: 5 },
AS: 'top_movies'
}
}]
});

assert.ok(reply !== null && typeof reply === 'object');
assert.ok(Array.isArray(reply.results));
assert.ok(reply.results.length > 0);
// Each group exposes the COLLECT alias holding an array of per-document maps
for (const group of reply.results) {
assert.ok('top_movies' in group);
assert.ok(Array.isArray((group as Record<string, unknown>).top_movies));
}
}, GLOBAL.SERVERS.OPEN_UNSTABLE);
});
48 changes: 46 additions & 2 deletions packages/search/lib/commands/AGGREGATE.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,8 @@ export const FT_AGGREGATE_GROUP_BY_REDUCERS = {
QUANTILE: 'QUANTILE',
TOLIST: 'TOLIST',
FIRST_VALUE: 'FIRST_VALUE',
RANDOM_SAMPLE: 'RANDOM_SAMPLE'
RANDOM_SAMPLE: 'RANDOM_SAMPLE',
COLLECT: 'COLLECT'
} as const;

type FT_AGGREGATE_GROUP_BY_REDUCERS = typeof FT_AGGREGATE_GROUP_BY_REDUCERS;
Expand Down Expand Up @@ -89,7 +90,17 @@ interface RandomSampleReducer extends GroupByReducerWithProperty<FT_AGGREGATE_GR
sampleSize: number;
}

export type GroupByReducers = CountReducer | CountDistinctReducer | CountDistinctishReducer | SumReducer | MinReducer | MaxReducer | AvgReducer | StdDevReducer | QuantileReducer | ToListReducer | FirstValueReducer | RandomSampleReducer;
interface CollectReducer extends GroupByReducer<FT_AGGREGATE_GROUP_BY_REDUCERS['COLLECT']> {
FIELDS: '*' | RediSearchProperty | Array<RediSearchProperty>;
DISTINCT?: boolean;
SORTBY?: SortByProperty | Array<SortByProperty>;
LIMIT?: {
from: number;
size: number;
};
}

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

interface GroupByStep extends AggregateStep<FT_AGGREGATE_STEPS['GROUPBY']> {
properties?: RediSearchProperty | Array<RediSearchProperty>;
Expand Down Expand Up @@ -365,6 +376,39 @@ export function parseGroupByReducer(parser: CommandParser, reducer: GroupByReduc
case FT_AGGREGATE_GROUP_BY_REDUCERS.RANDOM_SAMPLE:
parser.push('2', reducer.property, reducer.sampleSize.toString());
break;

case FT_AGGREGATE_GROUP_BY_REDUCERS.COLLECT: {
const args: Array<RedisArgument> = ['FIELDS'];

// Wildcard is a bare token (`FIELDS *`); an explicit list is count-prefixed
// (`FIELDS <num> field ...`).
if (reducer.FIELDS === '*') {
args.push('*');
} else {
const fields = Array.isArray(reducer.FIELDS) ? reducer.FIELDS : [reducer.FIELDS];
args.push(fields.length.toString(), ...fields);
}

if (reducer.DISTINCT) {
args.push('DISTINCT');
}

if (reducer.SORTBY) {
const sortBys = Array.isArray(reducer.SORTBY) ? reducer.SORTBY : [reducer.SORTBY];
const sortArgs: Array<RedisArgument> = [];
for (const sortBy of sortBys) {
pushSortByProperty(sortArgs, sortBy);
}
args.push('SORTBY', sortArgs.length.toString(), ...sortArgs);
}

if (reducer.LIMIT) {
args.push('LIMIT', reducer.LIMIT.from.toString(), reducer.LIMIT.size.toString());
}

parser.pushVariadicWithLength(args);
break;
}
}

if (reducer.AS) {
Expand Down
9 changes: 9 additions & 0 deletions packages/search/lib/test-utils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,15 @@ export const GLOBAL = {
ft: RediSearch
}
}
},
OPEN_UNSTABLE: {
serverArguments: ['--search-enable-unstable-features', 'yes'],
clientOptions: {
RESP: 3 as const,
modules: {
ft: RediSearch
}
}
}
}
};
Loading