Skip to content

Commit bceb8ab

Browse files
Hailong-amHailong-am
authored
Add auto aggregation suggest for t2v (#514)
* Add auto suggestion for t2v Signed-off-by: Hailong Cui <[email protected]> * add changelog Signed-off-by: Hailong Cui <[email protected]> * update class name Signed-off-by: Hailong Cui <[email protected]> * add null check Signed-off-by: Hailong Cui <[email protected]> --------- Signed-off-by: Hailong Cui <[email protected]>
1 parent c7d1c37 commit bceb8ab

File tree

5 files changed

+276
-10
lines changed

5 files changed

+276
-10
lines changed

‎CHANGELOG.md

+2
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,8 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
88

99
### Enhancements
1010

11+
- Add auto suggested aggregation for text2Viz ([#514](https://github.com/opensearch-project/dashboards-assistant/pull/514))
12+
1113
### Bug Fixes
1214

1315
### Infrastructure

‎public/components/visualization/text2viz.tsx

+3
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { useLocation, useParams } from 'react-router-dom';
2626
import { Pipeline } from '../../utils/pipeline/pipeline';
2727
import { Text2PPLTask } from '../../utils/pipeline/text_to_ppl_task';
2828
import { PPLSampleTask } from '../../utils/pipeline/ppl_sample_task';
29+
import { PPLAggsAutoSuggestTask } from '../../utils/pipeline/ppl_aggs_auto_suggest_task';
2930
import { SourceSelector } from './source_selector';
3031
import type { IndexPattern } from '../../../../../src/plugins/data/public';
3132
import { EmbeddableRenderer } from '../../../../../src/plugins/embeddable/public';
@@ -111,6 +112,7 @@ export const Text2Viz = () => {
111112
if (text2vegaRef.current === null) {
112113
text2vegaRef.current = new Pipeline([
113114
new Text2PPLTask(http),
115+
new PPLAggsAutoSuggestTask(data.search),
114116
new PPLSampleTask(data.search),
115117
new Text2VegaTask(http, savedObjects),
116118
]);
@@ -252,6 +254,7 @@ export const Text2Viz = () => {
252254
inputQuestion,
253255
inputInstruction,
254256
dataSourceId: indexPattern.dataSourceRef?.id,
257+
timeFiledName: indexPattern.timeFieldName,
255258
});
256259

257260
if (usageCollection) {

‎public/utils/pipeline/ppl_aggs_auto_suggest_task.test.ts

+197
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import { DataPublicPluginStart } from '../../../../../src/plugins/data/public';
7+
import { PPLAggsAutoSuggestTask } from './ppl_aggs_auto_suggest_task';
8+
9+
describe('PPLAutoSuggestTask', () => {
10+
let pplAutoSuggestTask: PPLAggsAutoSuggestTask;
11+
let mockSearchClient: DataPublicPluginStart['search'];
12+
13+
beforeEach(() => {
14+
// Mock the search client
15+
mockSearchClient = {
16+
search: jest.fn(),
17+
aggs: {
18+
calculateAutoTimeExpression: jest.fn(),
19+
},
20+
};
21+
pplAutoSuggestTask = new PPLAggsAutoSuggestTask(mockSearchClient);
22+
});
23+
24+
it('should return original input if PPL is empty', async () => {
25+
const input = {
26+
ppl: '',
27+
dataSourceId: 'test-source',
28+
};
29+
30+
const result = await pplAutoSuggestTask.execute(input);
31+
expect(result).toEqual(input);
32+
});
33+
34+
it('should return original input if PPL already has aggregation', async () => {
35+
const input = {
36+
ppl: 'source = test | stats count()',
37+
dataSourceId: 'test-source',
38+
};
39+
40+
const result = await pplAutoSuggestTask.execute(input);
41+
expect(result).toEqual(input);
42+
});
43+
44+
it('should add simple count aggregation when no time field is provided', async () => {
45+
const input = {
46+
ppl: 'source = test',
47+
dataSourceId: 'test-source',
48+
};
49+
50+
const expected = {
51+
ppl: 'source = test | stats count()',
52+
dataSourceId: 'test-source',
53+
};
54+
55+
const result = await pplAutoSuggestTask.execute(input);
56+
expect(result).toEqual(expected);
57+
});
58+
59+
it('should add time-based aggregation when time field is provided', async () => {
60+
const input = {
61+
ppl: 'source = test',
62+
dataSourceId: 'test-source',
63+
timeFiledName: 'timestamp',
64+
};
65+
66+
// Mock successful search response
67+
const mockResponse = {
68+
rawResponse: {
69+
total: 1,
70+
jsonData: [
71+
{
72+
min: '2023-01-01',
73+
max: '2023-12-31',
74+
},
75+
],
76+
},
77+
};
78+
79+
mockSearchClient.search.mockReturnValue({
80+
toPromise: () => Promise.resolve(mockResponse),
81+
});
82+
mockSearchClient.aggs.calculateAutoTimeExpression.mockReturnValue('1d');
83+
84+
const expected = {
85+
ppl: 'source = test | stats count() by span(timestamp, 1d)',
86+
dataSourceId: 'test-source',
87+
timeFiledName: 'timestamp',
88+
};
89+
90+
const result = await pplAutoSuggestTask.execute(input);
91+
expect(result).toEqual(expected);
92+
expect(mockSearchClient.search).toHaveBeenCalledWith(
93+
{
94+
params: {
95+
body: {
96+
query: 'source = test | stats min(timestamp) as min, max(timestamp) as max',
97+
},
98+
},
99+
dataSourceId: 'test-source',
100+
},
101+
{ strategy: 'pplraw' }
102+
);
103+
});
104+
105+
it('should handle empty search results for time-based queries', async () => {
106+
const input = {
107+
ppl: 'source = test',
108+
dataSourceId: 'test-source',
109+
timeFiledName: 'timestamp',
110+
};
111+
112+
// Mock empty search response
113+
const mockResponse = {
114+
rawResponse: {
115+
total: 0,
116+
jsonData: [],
117+
},
118+
};
119+
120+
mockSearchClient.search.mockReturnValue({
121+
toPromise: () => Promise.resolve(mockResponse),
122+
});
123+
124+
const expected = {
125+
ppl: 'source = test | stats count()',
126+
dataSourceId: 'test-source',
127+
timeFiledName: 'timestamp',
128+
};
129+
130+
const result = await pplAutoSuggestTask.execute(input);
131+
expect(result).toEqual(expected);
132+
});
133+
134+
it('should handle null search results for time-based queries', async () => {
135+
const input = {
136+
ppl: 'source = test',
137+
dataSourceId: 'test-source',
138+
timeFiledName: 'timestamp',
139+
};
140+
141+
// Mock null search response
142+
const mockResponse = null;
143+
144+
mockSearchClient.search.mockReturnValue({
145+
toPromise: () => Promise.resolve(mockResponse),
146+
});
147+
148+
const expected = {
149+
ppl: 'source = test | stats count()',
150+
dataSourceId: 'test-source',
151+
timeFiledName: 'timestamp',
152+
};
153+
154+
const result = await pplAutoSuggestTask.execute(input);
155+
expect(result).toEqual(expected);
156+
});
157+
158+
it('should handle search errors', async () => {
159+
const input = {
160+
ppl: 'source = test',
161+
dataSourceId: 'test-source',
162+
timeFiledName: 'timestamp',
163+
};
164+
165+
mockSearchClient.search.mockReturnValue({
166+
toPromise: () => Promise.reject(new Error('Search failed')),
167+
});
168+
169+
const expected = {
170+
...input,
171+
ppl: 'source = test | stats count()',
172+
};
173+
174+
const result = await pplAutoSuggestTask.execute(input);
175+
expect(result).toEqual(expected);
176+
});
177+
178+
describe('isPPLHasAggregation', () => {
179+
it('should detect stats command with various spacing', () => {
180+
const testCases = [
181+
'source = test | stats count()',
182+
'source = test |stats count()',
183+
'source = test| stats count()',
184+
'source = test|stats count()',
185+
];
186+
187+
testCases.forEach((ppl) => {
188+
expect(pplAutoSuggestTask.isPPLHasAggregation(ppl)).toBeTruthy();
189+
});
190+
});
191+
192+
it('should return false for queries without stats', () => {
193+
const ppl = 'source = test | where count > 0';
194+
expect(pplAutoSuggestTask.isPPLHasAggregation(ppl)).toBeFalsy();
195+
});
196+
});
197+
});

‎public/utils/pipeline/ppl_aggs_auto_suggest_task.ts

+73
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
import { Task } from './task';
7+
import { DataPublicPluginStart } from '../../../../../src/plugins/data/public';
8+
9+
interface Input {
10+
ppl: string;
11+
dataSourceId: string | undefined;
12+
timeFiledName?: string;
13+
}
14+
15+
export class PPLAggsAutoSuggestTask extends Task<Input, Input> {
16+
searchClient: DataPublicPluginStart['search'];
17+
18+
constructor(searchClient: DataPublicPluginStart['search']) {
19+
super();
20+
this.searchClient = searchClient;
21+
}
22+
23+
async execute<T extends Input>(v: T) {
24+
let ppl = v.ppl;
25+
26+
if (ppl) {
27+
const isPPLHasAgg = this.isPPLHasAggregation(ppl);
28+
// if no aggregation, will try auto suggest one
29+
if (!isPPLHasAgg) {
30+
if (v.timeFiledName) {
31+
const dateRangePPL = `${ppl} | stats min(${v.timeFiledName}) as min, max(${v.timeFiledName}) as max`;
32+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
33+
let res: any;
34+
try {
35+
res = await this.searchClient
36+
.search(
37+
{ params: { body: { query: dateRangePPL } }, dataSourceId: v.dataSourceId },
38+
{ strategy: 'pplraw' }
39+
)
40+
// eslint-disable-next-line @typescript-eslint/no-explicit-any
41+
.toPromise<any>();
42+
43+
if (!res || !res.rawResponse || res.rawResponse.total === 0) {
44+
ppl = `${ppl} | stats count()`;
45+
} else {
46+
// get min and max and calculate proper interval by range
47+
const min = res.rawResponse.jsonData[0].min;
48+
const max = res.rawResponse.jsonData[0].max;
49+
const interval =
50+
this.searchClient.aggs.calculateAutoTimeExpression({ from: min, to: max }) || '1d';
51+
52+
ppl = `${ppl} | stats count() by span(${v.timeFiledName}, ${interval})`;
53+
}
54+
} catch (e) {
55+
ppl = `${ppl} | stats count()`;
56+
}
57+
} else {
58+
ppl = `${ppl} | stats count()`;
59+
}
60+
61+
// override the original input with suggested ppl
62+
return { ...v, ppl };
63+
}
64+
}
65+
// directly return original input
66+
return v;
67+
}
68+
69+
private isPPLHasAggregation(ppl: string) {
70+
const statsRegex = /\|\s*stats\s+/i;
71+
return statsRegex.test(ppl);
72+
}
73+
}

‎public/utils/pipeline/text_to_ppl_task.ts

+1-10
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ interface Input {
1919
inputQuestion: string;
2020
index: string;
2121
dataSourceId?: string;
22+
timeFiledName?: string;
2223
}
2324

2425
export class Text2PPLTask extends Task<Input, Input & { ppl: string }> {
@@ -39,16 +40,6 @@ export class Text2PPLTask extends Task<Input, Input & { ppl: string }> {
3940
);
4041
}
4142

42-
if (ppl) {
43-
const statsRegex = /\|\s*stats\s+/i;
44-
const hasStats = statsRegex.test(ppl);
45-
if (!hasStats) {
46-
throw new Error(
47-
`The generated PPL query doesn't seem to contain an aggregation. Ensure your question contains an aggregation (e.g., average, sum, or count) to create a meaningful visualization. Generated PPL: ${ppl}`
48-
);
49-
}
50-
}
51-
5243
return { ...v, ppl };
5344
}
5445

0 commit comments

Comments
 (0)