Skip to content

Commit 1c7e72e

Browse files
pulpdrewpulpdrew
committed
feat: Improve MV source column inference
1 parent 77a1aa9 commit 1c7e72e

2 files changed

Lines changed: 241 additions & 34 deletions

File tree

‎packages/app/src/utils/__tests__/materializedViews.test.ts‎

Lines changed: 169 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import {
88
import { getMetadata } from '@/metadata';
99

1010
import {
11+
getSourceTableColumn,
1112
inferMaterializedViewConfig,
1213
inferTimestampColumnGranularity,
1314
parseSummedColumns,
@@ -87,20 +88,24 @@ describe('inferMaterializedViewConfig', () => {
8788
name: 'SpanKind',
8889
type: 'LowCardinality(String)',
8990
}),
91+
createMockColumnMeta({
92+
name: 'quantileDuration',
93+
type: 'AggregateFunction(quantile(0.5), UInt64)',
94+
}),
9095
createMockColumnMeta({
9196
name: 'count',
9297
type: 'UInt64',
9398
}),
9499
createMockColumnMeta({
95-
name: 'sum__Duration',
100+
name: 'sumDuration',
96101
type: 'UInt64',
97102
}),
98103
] as ColumnMeta[],
99104
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
100105
meta: {
101106
engine: 'SummingMergeTree',
102107
engine_full:
103-
'SummingMergeTree((count, sum__Duration)) ORDER BY (Timestamp, ServiceName, SpanKind) SETTINGS index_granularity = 8192',
108+
'SummingMergeTree((count, sumDuration)) ORDER BY (Timestamp, ServiceName, SpanKind) SETTINGS index_granularity = 8192',
104109
database: 'test_db',
105110
name: 'test_mv_target_table_summing',
106111
primary_key: 'Timestamp, ServiceName, SpanKind',
@@ -131,11 +136,27 @@ describe('inferMaterializedViewConfig', () => {
131136
name: 'test_mv',
132137
create_table_query: `CREATE MATERIALIZED VIEW test_db.test_mv TO test_db.test_mv_target_table AS
133138
SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind,
134-
count(*) AS count, sum(Duration) AS sum__Duration, quantileState(0.5)(Duration) AS quantile__Duration
139+
count(*) AS count, sum(Duration) AS sum__Duration, histogram(20)(Duration) as histogram__Duration, quantileState(0.5)(Duration) AS quantile__Duration
135140
FROM test_source_table
136141
GROUP BY Timestamp, ServiceName, SpanKind`,
137142
as_select: `SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind,
138-
count(*) AS count, sum(Duration) AS sum__Duration, quantileState(0.5)(Duration) AS quantile__Duration
143+
count(*) AS count, sum(Duration) AS sum__Duration, histogram(20)(Duration) as histogram__Duration, quantileState(0.5)(Duration) AS quantile__Duration
144+
FROM test_source_table
145+
GROUP BY Timestamp, ServiceName, SpanKind`,
146+
} as unknown as TableMetadata,
147+
};
148+
149+
const summingMergeTreeMV = {
150+
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
151+
meta: {
152+
engine: 'MaterializedView',
153+
database: 'test_db',
154+
name: 'test_mv_summing',
155+
create_table_query: `CREATE MATERIALIZED VIEW test_db.test_mv_summing TO test_db.test_mv_target_table_summing AS
156+
SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind, count() AS count, sum(Duration) AS sumDuration, quantileState(0.5)(Duration) AS quantileDuration
157+
FROM test_source_table
158+
GROUP BY Timestamp, ServiceName, SpanKind`,
159+
as_select: `SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind, count() AS count, sum(Duration) AS sumDuration, quantileState(0.5)(Duration) AS quantileDuration
139160
FROM test_source_table
140161
GROUP BY Timestamp, ServiceName, SpanKind`,
141162
} as unknown as TableMetadata,
@@ -148,6 +169,8 @@ describe('inferMaterializedViewConfig', () => {
148169
.mockImplementation(({ tableName }) => {
149170
if (tableName === 'test_mv') {
150171
return Promise.resolve(mv.meta);
172+
} else if (tableName === 'test_mv_summing') {
173+
return Promise.resolve(summingMergeTreeMV.meta);
151174
} else if (tableName === 'test_mv_target_table') {
152175
return Promise.resolve(mvTargetTable.meta);
153176
} else if (tableName === 'test_mv_target_table_summing') {
@@ -170,7 +193,17 @@ describe('inferMaterializedViewConfig', () => {
170193

171194
mockMetadata.queryMaterializedViewsByTarget = jest
172195
.fn()
173-
.mockResolvedValue([{ tableName: 'test_mv', databaseName: 'test_db' }]);
196+
.mockImplementation(({ tableName }) => {
197+
return Promise.resolve([
198+
{
199+
databaseName: 'test_db',
200+
tableName:
201+
tableName === 'test_mv_target_table_summing'
202+
? 'test_mv_summing'
203+
: 'test_mv',
204+
},
205+
]);
206+
});
174207
});
175208

176209
afterEach(() => {
@@ -300,14 +333,19 @@ describe('inferMaterializedViewConfig', () => {
300333
timestampColumn: 'Timestamp',
301334
minGranularity: '1 hour',
302335
aggregatedColumns: [
336+
{
337+
aggFn: 'quantile',
338+
mvColumn: 'quantileDuration',
339+
sourceColumn: 'Duration',
340+
},
303341
{
304342
aggFn: 'count',
305343
mvColumn: 'count',
306344
sourceColumn: '',
307345
},
308346
{
309347
aggFn: 'sum',
310-
mvColumn: 'sum__Duration',
348+
mvColumn: 'sumDuration',
311349
sourceColumn: 'Duration',
312350
},
313351
],
@@ -533,3 +571,128 @@ describe('parseSummedColumns', () => {
533571
expect(parsed).toEqual(new Set(['count', 'sum__Duration']));
534572
});
535573
});
574+
575+
describe('getSourceTableColumn', () => {
576+
it('should return empty string if no matching source column is found', () => {
577+
const sourceTableColumns: ColumnMeta[] = [
578+
createMockColumnMeta({ name: 'Duration', type: 'UInt64' }),
579+
createMockColumnMeta({ name: 'Value', type: 'UInt64' }),
580+
];
581+
582+
const targetTableColumn = createMockColumnMeta({
583+
name: 'sum__NonExistentColumn',
584+
type: 'SimpleAggregateFunction(sum, UInt64)',
585+
});
586+
const sourceColumn = getSourceTableColumn(
587+
'sum',
588+
targetTableColumn,
589+
sourceTableColumns,
590+
);
591+
expect(sourceColumn).toBe('');
592+
});
593+
594+
it('should return empty string if the aggFn is count', () => {
595+
const sourceTableColumns: ColumnMeta[] = [
596+
createMockColumnMeta({ name: 'Duration', type: 'UInt64' }),
597+
createMockColumnMeta({ name: 'Value', type: 'UInt64' }),
598+
];
599+
600+
const targetTableColumn = createMockColumnMeta({
601+
name: 'count',
602+
type: 'SimpleAggregateFunction(count, UInt64)',
603+
});
604+
const sourceColumn = getSourceTableColumn(
605+
'count',
606+
targetTableColumn,
607+
sourceTableColumns,
608+
);
609+
expect(sourceColumn).toBe('');
610+
});
611+
612+
it('should match source columns based on convention', () => {
613+
const sourceTableColumns: ColumnMeta[] = [
614+
createMockColumnMeta({ name: 'Duration', type: 'UInt64' }),
615+
createMockColumnMeta({ name: 'Value', type: 'UInt64' }),
616+
];
617+
618+
const targetTableColumnSum = createMockColumnMeta({
619+
name: 'sum__Duration',
620+
type: 'SimpleAggregateFunction(sum, UInt64)',
621+
});
622+
const sourceColumnForSum = getSourceTableColumn(
623+
'sum',
624+
targetTableColumnSum,
625+
sourceTableColumns,
626+
);
627+
expect(sourceColumnForSum).toBe('Duration');
628+
});
629+
630+
it('should match source column based on MV DDL expressions', () => {
631+
const sourceTableColumns: ColumnMeta[] = [
632+
createMockColumnMeta({ name: 'Duration', type: 'UInt64' }),
633+
createMockColumnMeta({ name: 'Value', type: 'UInt64' }),
634+
];
635+
636+
const targetTableColumnQuantile = createMockColumnMeta({
637+
name: 'quantileDuration',
638+
type: 'AggregateFunction(quantile(0.5), UInt64)',
639+
});
640+
641+
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
642+
const mvMetadata: TableMetadata = {
643+
as_select:
644+
'SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind, count() AS count, sum(Duration) AS sumDuration, quantileState(0.5)(Duration) AS quantileDuration FROM test_source_table GROUP BY Timestamp, ServiceName, SpanKind',
645+
} as unknown as TableMetadata;
646+
647+
const sourceColumnForQuantile = getSourceTableColumn(
648+
'quantile',
649+
targetTableColumnQuantile,
650+
sourceTableColumns,
651+
mvMetadata,
652+
);
653+
654+
expect(sourceColumnForQuantile).toBe('Duration');
655+
});
656+
657+
it('should match source column based on MV DDL expressions when there are overlapping source column names', () => {
658+
const sourceTableColumns: ColumnMeta[] = [
659+
createMockColumnMeta({ name: 'MaxDuration', type: 'UInt64' }),
660+
createMockColumnMeta({ name: 'Duration', type: 'UInt64' }),
661+
createMockColumnMeta({ name: 'Value', type: 'UInt64' }),
662+
];
663+
664+
const targetTableColumnMaxMax = createMockColumnMeta({
665+
name: 'maxMaxDuration',
666+
type: 'AggregateFunction(max, UInt64)',
667+
});
668+
669+
const targetTableColumnMax = createMockColumnMeta({
670+
name: 'maxDuration',
671+
type: 'AggregateFunction(max, UInt64)',
672+
});
673+
674+
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
675+
const mvMetadata: TableMetadata = {
676+
as_select:
677+
'SELECT toStartOfHour(Timestamp) AS Timestamp, ServiceName, SpanKind, count() AS count, max(MaxDuration) AS maxMaxDuration, max(Duration) AS maxDuration FROM test_source_table GROUP BY Timestamp, ServiceName, SpanKind',
678+
} as unknown as TableMetadata;
679+
680+
const sourceColumnForMax = getSourceTableColumn(
681+
'max',
682+
targetTableColumnMax,
683+
sourceTableColumns,
684+
mvMetadata,
685+
);
686+
687+
expect(sourceColumnForMax).toBe('Duration');
688+
689+
const sourceColumnForMaxMax = getSourceTableColumn(
690+
'max',
691+
targetTableColumnMaxMax,
692+
sourceTableColumns,
693+
mvMetadata,
694+
);
695+
696+
expect(sourceColumnForMaxMax).toBe('MaxDuration');
697+
});
698+
});

‎packages/app/src/utils/materializedViews.ts‎

Lines changed: 72 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import {
2+
ColumnMeta,
23
extractColumnReferencesFromKey,
34
filterColumnMetaByType,
45
JSDataType,
@@ -247,6 +248,7 @@ export function parseSummedColumns(mvTableMetadata: TableMetadata) {
247248
}
248249

249250
// Extract the column list from the engine parameters
251+
// SummingMergeTree(col1) or SummingMergeTree((col1, col2, ...))
250252
const engineParamStr = mvTableMetadata.engine_full?.match(
251253
/SummingMergeTree\((\(?[^(]*)\)/,
252254
)?.[1];
@@ -262,20 +264,53 @@ export function parseSummedColumns(mvTableMetadata: TableMetadata) {
262264
}
263265
}
264266

265-
function getSourceTableColumn(
266-
mvColumnName: string,
267+
function escapeRegExp(s: string) {
268+
// $& means the whole matched string
269+
return s.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
270+
}
271+
272+
export function getSourceTableColumn(
267273
aggFn: string,
268-
sourceTableColumnNames: Set<string>,
274+
targetTableColumn: ColumnMeta,
275+
sourceTableColumns: ColumnMeta[],
276+
mvMetadata?: TableMetadata,
269277
) {
278+
if (aggFn === 'count') {
279+
// Count may not have a source column
280+
return '';
281+
}
282+
270283
// By convention: MV Columns are named "<aggFn>__<sourceColumn>"
271-
const nameSuffix = mvColumnName.split('__')[1];
272-
return sourceTableColumnNames.has(nameSuffix) && aggFn !== 'count'
273-
? nameSuffix
274-
: '';
284+
const nameSuffix = targetTableColumn.name.split('__')[1];
285+
if (sourceTableColumns.find(col => col.name === nameSuffix)) {
286+
return nameSuffix;
287+
}
288+
289+
// Try to infer from the MV's SELECT expression
290+
if (mvMetadata) {
291+
const selectExpressions = extractSelectExpressions(mvMetadata);
292+
const matchingSelectExpression = selectExpressions.find(expr =>
293+
// Use endsWith because the expression must have an alias
294+
// matching the target column name.
295+
expr.endsWith(targetTableColumn.name),
296+
);
297+
const matchingSourceColumn =
298+
matchingSelectExpression &&
299+
sourceTableColumns.find(col =>
300+
new RegExp(`\\b${escapeRegExp(col.name)}\\b`).test(
301+
matchingSelectExpression,
302+
),
303+
);
304+
if (matchingSourceColumn) {
305+
return matchingSourceColumn.name;
306+
}
307+
}
308+
309+
return '';
275310
}
276311

277312
/**
278-
* Attempts to a MaterializedViewConfiguration object from the given TableConnections
313+
* Attempts to create a MaterializedViewConfiguration object from the given TableConnections
279314
* by introspecting the view, target table, and source table.
280315
*
281316
* @param mvTableOrView - A TableConnection representing either the materialized view or the target table.
@@ -321,18 +356,19 @@ export async function inferMaterializedViewConfig(
321356
}),
322357
]);
323358

324-
const sourceTableColumnNames = new Set(
325-
sourceTableColumns.map(col => col.name),
326-
);
327-
328359
const aggregatedColumns: MaterializedViewConfiguration['aggregatedColumns'] =
329360
mvTableColumns
330-
.filter(col => col.type.includes('AggregateFunction'))
331-
.map(col => {
332-
let aggFn: string | undefined = col.type.match(
361+
.filter(targetTableColumn =>
362+
targetTableColumn.type.includes('AggregateFunction'),
363+
)
364+
.map(targetTableColumn => {
365+
let aggFn: string | undefined = targetTableColumn.type.match(
333366
/AggregateFunction\(([a-zA-Z0-9_]+)/,
334367
)?.[1];
335-
if (aggFn === 'sum' && col.name.toLowerCase().includes('count')) {
368+
if (
369+
aggFn === 'sum' &&
370+
targetTableColumn.name.toLowerCase().includes('count')
371+
) {
336372
aggFn = 'count';
337373
} else if (aggFn?.startsWith('quantile')) {
338374
aggFn = 'quantile';
@@ -343,13 +379,14 @@ export async function inferMaterializedViewConfig(
343379
}
344380

345381
const sourceColumn = getSourceTableColumn(
346-
col.name,
347382
aggFn,
348-
sourceTableColumnNames,
383+
targetTableColumn,
384+
sourceTableColumns,
385+
mvMetadata,
349386
);
350387

351388
return {
352-
mvColumn: col.name,
389+
mvColumn: targetTableColumn.name,
353390
aggFn,
354391
sourceColumn,
355392
};
@@ -367,16 +404,23 @@ export async function inferMaterializedViewConfig(
367404
? 'count'
368405
: 'sum';
369406

370-
const sourceColumn = getSourceTableColumn(
371-
summedColumn,
372-
aggFn,
373-
sourceTableColumnNames,
407+
const summedColumnMeta = mvTableColumns.find(
408+
col => col.name === summedColumn,
374409
);
375-
aggregatedColumns.push({
376-
mvColumn: summedColumn,
377-
aggFn,
378-
sourceColumn,
379-
});
410+
411+
if (summedColumnMeta) {
412+
const sourceColumn = getSourceTableColumn(
413+
aggFn,
414+
summedColumnMeta,
415+
sourceTableColumns,
416+
mvMetadata,
417+
);
418+
aggregatedColumns.push({
419+
mvColumn: summedColumn,
420+
aggFn,
421+
sourceColumn,
422+
});
423+
}
380424
}
381425

382426
// Infer the timestamp column

0 commit comments

Comments
 (0)