diff --git a/.changeset/tiny-forks-deny.md b/.changeset/tiny-forks-deny.md new file mode 100644 index 0000000000..0a5a7a9eec --- /dev/null +++ b/.changeset/tiny-forks-deny.md @@ -0,0 +1,7 @@ +--- +'@hyperdx/common-utils': patch +'@hyperdx/api': patch +'@hyperdx/app': patch +--- + +feat: support sample-weighted aggregations for sampled trace data diff --git a/.gitignore b/.gitignore index dd26f4fac5..e22532cf59 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,9 @@ packages/app/next-env.d.ts # optional npm cache directory **/.npm +# npm lockfile (project uses yarn) +package-lock.json + # dependency directories **/node_modules diff --git a/docker/clickhouse/local/init-db-e2e.sh b/docker/clickhouse/local/init-db-e2e.sh index 4166a6e3e4..7b6fe87889 100755 --- a/docker/clickhouse/local/init-db-e2e.sh +++ b/docker/clickhouse/local/init-db-e2e.sh @@ -81,6 +81,7 @@ CREATE TABLE IF NOT EXISTS ${DATABASE}.e2e_otel_traces \`Links.TraceState\` Array(String) CODEC(ZSTD(1)), \`Links.Attributes\` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), \`__hdx_materialized_rum.sessionId\` String MATERIALIZED ResourceAttributes['rum.sessionId'] CODEC(ZSTD(1)), + \`SampleRate\` UInt64 MATERIALIZED greatest(toUInt64OrZero(SpanAttributes['SampleRate']), 1) CODEC(T64, ZSTD(1)), INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1, INDEX idx_rum_session_id __hdx_materialized_rum.sessionId TYPE bloom_filter(0.001) GRANULARITY 1, INDEX idx_res_attr_key mapKeys(ResourceAttributes) TYPE bloom_filter(0.01) GRANULARITY 1, diff --git a/docker/otel-collector/schema/seed/00005_otel_traces.sql b/docker/otel-collector/schema/seed/00005_otel_traces.sql index 98341a6733..853ea471fe 100644 --- a/docker/otel-collector/schema/seed/00005_otel_traces.sql +++ b/docker/otel-collector/schema/seed/00005_otel_traces.sql @@ -24,6 +24,7 @@ CREATE TABLE IF NOT EXISTS ${DATABASE}.otel_traces `Links.TraceState` Array(String) CODEC(ZSTD(1)), `Links.Attributes` Array(Map(LowCardinality(String), String)) CODEC(ZSTD(1)), `__hdx_materialized_rum.sessionId` String MATERIALIZED ResourceAttributes['rum.sessionId'] CODEC(ZSTD(1)), + `SampleRate` UInt64 MATERIALIZED greatest(toUInt64OrZero(SpanAttributes['SampleRate']), 1) CODEC(T64, ZSTD(1)), INDEX idx_trace_id TraceId TYPE bloom_filter(0.001) GRANULARITY 1, INDEX idx_rum_session_id __hdx_materialized_rum.sessionId TYPE bloom_filter(0.001) GRANULARITY 1, INDEX idx_res_attr_key mapKeys(ResourceAttributes) TYPE bloom_filter(0.01) GRANULARITY 1, diff --git a/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.down.sql b/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.down.sql new file mode 100644 index 0000000000..c21fb812e8 --- /dev/null +++ b/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.down.sql @@ -0,0 +1 @@ +ALTER TABLE ${DATABASE}.otel_traces DROP COLUMN IF EXISTS `SampleRate`; diff --git a/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.up.sql b/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.up.sql new file mode 100644 index 0000000000..5f2cf8cfa5 --- /dev/null +++ b/packages/api/migrations/ch/000002_add_sample_rate_column_to_otel_traces.up.sql @@ -0,0 +1,4 @@ +ALTER TABLE ${DATABASE}.otel_traces + ADD COLUMN IF NOT EXISTS `SampleRate` UInt64 + MATERIALIZED greatest(toUInt64OrZero(SpanAttributes['SampleRate']), 1) + CODEC(T64, ZSTD(1)); diff --git a/packages/api/src/controllers/ai.ts b/packages/api/src/controllers/ai.ts index 2dced86237..40c1bec8a4 100644 --- a/packages/api/src/controllers/ai.ts +++ b/packages/api/src/controllers/ai.ts @@ -8,6 +8,7 @@ import { AILineTableResponse, AssistantLineTableConfigSchema, ChartConfigWithDateRange, + SourceKind, } from '@hyperdx/common-utils/dist/types'; import type { LanguageModel } from 'ai'; import * as chrono from 'chrono-node'; @@ -328,6 +329,11 @@ export function getChartConfigFromResolvedConfig( connection: source.connection.toString(), groupBy: resObject.groupBy, timestampValueExpression: source.timestampValueExpression, + ...(source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), dateRange: [dateRange[0].toString(), dateRange[1].toString()], markdown: resObject.markdown, granularity: 'auto', diff --git a/packages/api/src/models/source.ts b/packages/api/src/models/source.ts index 1e6fb20830..25201dd0f9 100644 --- a/packages/api/src/models/source.ts +++ b/packages/api/src/models/source.ts @@ -163,6 +163,7 @@ export const TraceSource = Source.discriminator( parentSpanIdExpression: String, spanNameExpression: String, spanKindExpression: String, + sampleRateExpression: String, logSourceId: String, sessionSourceId: String, metricSourceId: String, diff --git a/packages/api/src/routers/external-api/v2/charts.ts b/packages/api/src/routers/external-api/v2/charts.ts index 7f574ed9ce..33b2022c6a 100644 --- a/packages/api/src/routers/external-api/v2/charts.ts +++ b/packages/api/src/routers/external-api/v2/charts.ts @@ -280,6 +280,11 @@ const buildChartConfigFromRequest = async ( ], where: '', timestampValueExpression: source.timestampValueExpression, + ...(source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), dateRange: [new Date(params.startTime), new Date(params.endTime)] as [ Date, Date, diff --git a/packages/api/src/tasks/checkAlerts/index.ts b/packages/api/src/tasks/checkAlerts/index.ts index d69fbd4023..7beda653f7 100644 --- a/packages/api/src/tasks/checkAlerts/index.ts +++ b/packages/api/src/tasks/checkAlerts/index.ts @@ -107,6 +107,11 @@ export async function computeAliasWithClauses( source.kind === SourceKind.Log || source.kind === SourceKind.Trace ? source.implicitColumnExpression : undefined, + ...(source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), timestampValueExpression: source.timestampValueExpression, }; const query = await renderChartConfig(config, metadata, source.querySettings); @@ -453,6 +458,11 @@ const getChartConfigFromAlert = ( source.kind === SourceKind.Log || source.kind === SourceKind.Trace ? source.implicitColumnExpression : undefined, + ...(source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), timestampValueExpression: source.timestampValueExpression, }; } else if (details.taskType === AlertTaskType.TILE) { @@ -474,6 +484,12 @@ const getChartConfigFromAlert = ( source.kind === SourceKind.Log || source.kind === SourceKind.Trace ? source.implicitColumnExpression : undefined; + const sampleWeightExpression = + source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression + ? source.sampleRateExpression + : undefined; const metricTables = source.kind === SourceKind.Metric ? source.metricTables : undefined; return { @@ -486,6 +502,7 @@ const getChartConfigFromAlert = ( granularity: `${windowSizeInMins} minute`, groupBy: tile.config.groupBy, implicitColumnExpression, + sampleWeightExpression, metricTables, select: tile.config.select, timestampValueExpression: source.timestampValueExpression, diff --git a/packages/api/src/tasks/checkAlerts/template.ts b/packages/api/src/tasks/checkAlerts/template.ts index dce38c276c..2e6663645e 100644 --- a/packages/api/src/tasks/checkAlerts/template.ts +++ b/packages/api/src/tasks/checkAlerts/template.ts @@ -598,6 +598,11 @@ ${targetTemplate}`; where: savedSearch.where, whereLanguage: savedSearch.whereLanguage, implicitColumnExpression: source.implicitColumnExpression, + ...(source.kind === SourceKind.Trace && + 'sampleRateExpression' in source && + source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), timestampValueExpression: source.timestampValueExpression, orderBy: savedSearch.orderBy, limit: { diff --git a/packages/app/src/DBDashboardPage.tsx b/packages/app/src/DBDashboardPage.tsx index 9781278051..6ecc3ddd21 100644 --- a/packages/app/src/DBDashboardPage.tsx +++ b/packages/app/src/DBDashboardPage.tsx @@ -231,6 +231,9 @@ const Tile = forwardRef( ...chart.config, // Populate these two columns from the source to support Lucene-based filters ...pick(source, ['implicitColumnExpression', 'from']), + sampleWeightExpression: isTraceSource(source) + ? source.sampleRateExpression + : undefined, dateRange, granularity, filters, @@ -265,6 +268,9 @@ const Tile = forwardRef( isLogSource(source) || isTraceSource(source) ? source.implicitColumnExpression : undefined, + sampleWeightExpression: isTraceSource(source) + ? source.sampleRateExpression + : undefined, filters, metricTables: isMetricSource ? source.metricTables : undefined, }); diff --git a/packages/app/src/DBSearchPage.tsx b/packages/app/src/DBSearchPage.tsx index 454a9bded2..5719ae4758 100644 --- a/packages/app/src/DBSearchPage.tsx +++ b/packages/app/src/DBSearchPage.tsx @@ -638,6 +638,9 @@ function useSearchedConfigToChartConfig( whereLanguage: whereLanguage ?? 'sql', timestampValueExpression: sourceObj.timestampValueExpression, implicitColumnExpression: sourceObj.implicitColumnExpression, + ...(isTraceSource(sourceObj) && sourceObj.sampleRateExpression + ? { sampleWeightExpression: sourceObj.sampleRateExpression } + : {}), connection: sourceObj.connection, displayType: DisplayType.Search, orderBy: orderBy || defaultSearchConfig?.orderBy || defaultOrderBy, diff --git a/packages/app/src/ServicesDashboardPage.tsx b/packages/app/src/ServicesDashboardPage.tsx index 5795006043..ba876cf048 100644 --- a/packages/app/src/ServicesDashboardPage.tsx +++ b/packages/app/src/ServicesDashboardPage.tsx @@ -33,6 +33,9 @@ function pickSourceConfigFields(source: TSource) { ...(isLogSource(source) || isTraceSource(source) ? { implicitColumnExpression: source.implicitColumnExpression } : {}), + ...(isTraceSource(source) && source.sampleRateExpression + ? { sampleWeightExpression: source.sampleRateExpression } + : {}), }; } import { diff --git a/packages/app/src/SessionSubpanel.tsx b/packages/app/src/SessionSubpanel.tsx index 6ccf93d510..30cff55ccd 100644 --- a/packages/app/src/SessionSubpanel.tsx +++ b/packages/app/src/SessionSubpanel.tsx @@ -188,6 +188,9 @@ function useSessionChartConfigs({ where, timestampValueExpression: traceSource.timestampValueExpression, implicitColumnExpression: traceSource.implicitColumnExpression, + ...(traceSource.sampleRateExpression && { + sampleWeightExpression: traceSource.sampleRateExpression, + }), connection: traceSource.connection, orderBy: `${traceSource.timestampValueExpression} ASC`, limit: { diff --git a/packages/app/src/components/AlertPreviewChart.tsx b/packages/app/src/components/AlertPreviewChart.tsx index b1baabef37..4b96949b1e 100644 --- a/packages/app/src/components/AlertPreviewChart.tsx +++ b/packages/app/src/components/AlertPreviewChart.tsx @@ -74,6 +74,9 @@ export const AlertPreviewChart = ({ isLogSource(source) || isTraceSource(source) ? source.implicitColumnExpression : undefined, + sampleWeightExpression: isTraceSource(source) + ? source.sampleRateExpression + : undefined, groupBy, with: aliasWith, select: [ diff --git a/packages/app/src/components/ChartEditor/utils.ts b/packages/app/src/components/ChartEditor/utils.ts index e3e5e047a4..92ffaf243a 100644 --- a/packages/app/src/components/ChartEditor/utils.ts +++ b/packages/app/src/components/ChartEditor/utils.ts @@ -142,6 +142,9 @@ export function convertFormStateToChartConfig( isLogSource(source) || isTraceSource(source) ? source.implicitColumnExpression : undefined, + sampleWeightExpression: isTraceSource(source) + ? source.sampleRateExpression + : undefined, metricTables: isMetricSource(source) ? source.metricTables : undefined, where: form.where ?? '', select: isSelectEmpty diff --git a/packages/app/src/components/ServiceDashboardDbQuerySidePanel.tsx b/packages/app/src/components/ServiceDashboardDbQuerySidePanel.tsx index 8580e54fdf..d9621c4944 100644 --- a/packages/app/src/components/ServiceDashboardDbQuerySidePanel.tsx +++ b/packages/app/src/components/ServiceDashboardDbQuerySidePanel.tsx @@ -109,6 +109,9 @@ export default function ServiceDashboardDbQuerySidePanel({ 'connection', 'from', ]), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ @@ -146,6 +149,9 @@ export default function ServiceDashboardDbQuerySidePanel({ 'connection', 'from', ]), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ diff --git a/packages/app/src/components/ServiceDashboardEndpointPerformanceChart.tsx b/packages/app/src/components/ServiceDashboardEndpointPerformanceChart.tsx index 995ffa4ed3..f7bf988e2a 100644 --- a/packages/app/src/components/ServiceDashboardEndpointPerformanceChart.tsx +++ b/packages/app/src/components/ServiceDashboardEndpointPerformanceChart.tsx @@ -95,6 +95,9 @@ export default function ServiceDashboardEndpointPerformanceChart({ config={{ source: source.id, ...pick(source, ['timestampValueExpression', 'connection', 'from']), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ diff --git a/packages/app/src/components/ServiceDashboardEndpointSidePanel.tsx b/packages/app/src/components/ServiceDashboardEndpointSidePanel.tsx index 26998b5386..13824aa45e 100644 --- a/packages/app/src/components/ServiceDashboardEndpointSidePanel.tsx +++ b/packages/app/src/components/ServiceDashboardEndpointSidePanel.tsx @@ -116,6 +116,9 @@ export default function ServiceDashboardEndpointSidePanel({ 'connection', 'from', ]), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ @@ -159,6 +162,9 @@ export default function ServiceDashboardEndpointSidePanel({ 'connection', 'from', ]), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ diff --git a/packages/app/src/components/ServiceDashboardSlowestEventsTile.tsx b/packages/app/src/components/ServiceDashboardSlowestEventsTile.tsx index 048a3975c4..3dc76bc647 100644 --- a/packages/app/src/components/ServiceDashboardSlowestEventsTile.tsx +++ b/packages/app/src/components/ServiceDashboardSlowestEventsTile.tsx @@ -33,6 +33,9 @@ export default function SlowestEventsTile({ { source: source.id, ...pick(source, ['timestampValueExpression', 'connection', 'from']), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ @@ -117,6 +120,9 @@ export default function SlowestEventsTile({ 'connection', 'from', ]), + ...(source.sampleRateExpression && { + sampleWeightExpression: source.sampleRateExpression, + }), where: '', whereLanguage: 'sql', select: [ diff --git a/packages/app/src/components/Sources/SourceForm.tsx b/packages/app/src/components/Sources/SourceForm.tsx index 26d3b310c2..2a249f940c 100644 --- a/packages/app/src/components/Sources/SourceForm.tsx +++ b/packages/app/src/components/Sources/SourceForm.tsx @@ -1545,6 +1545,21 @@ export function TraceTableModelForm(props: TableModelProps) { placeholder="SpanAttributes" /> + + + = fromUnixTimestamp64Milli(1739318400000) AND __hdx_time_bucket2 <= fromUnixTimestamp64Milli(1765670400000)) GROUP BY toStartOfInterval(toDateTime(__hdx_time_bucket2), INTERVAL 1 minute) AS \`__hdx_time_bucket\` ORDER BY toStartOfInterval(toDateTime(__hdx_time_bucket2), INTERVAL 1 minute) AS \`__hdx_time_bucket\` LIMIT 10 SETTINGS short_circuit_function_evaluation = 'force_enable', optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000" `; +exports[`renderChartConfig sample-weighted aggregations should handle complex sampleWeightExpression like SpanAttributes map access 1`] = `"SELECT sum(greatest(toUInt64OrZero(toString(SpanAttributes['SampleRate'])), 1)) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should handle mixed weighted and passthrough aggregations 1`] = ` +"SELECT sum(greatest(toUInt64OrZero(toString(SampleRate)), 1)) AS \\"weighted_count\\",sumIf(toFloat64OrDefault(toString(Duration)) * greatest(toUInt64OrZero(toString(SampleRate)), 1), toFloat64OrDefault(toString(Duration)) IS NOT NULL) / sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1), toFloat64OrDefault(toString(Duration)) IS NOT NULL) AS \\"weighted_avg\\",min( + toFloat64OrDefault(toString(Duration)) + ) AS \\"min_duration\\",count(DISTINCT TraceId) AS \\"unique_traces\\" FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000" +`; + +exports[`renderChartConfig sample-weighted aggregations should leave count_distinct unchanged with sampleWeightExpression 1`] = `"SELECT count(DISTINCT TraceId) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should leave min/max unchanged with sampleWeightExpression 1`] = ` +"SELECT min( + toFloat64OrDefault(toString(Duration)) + ),max( + toFloat64OrDefault(toString(Duration)) + ) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000" +`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite avg to weighted average 1`] = `"SELECT sumIf(toFloat64OrDefault(toString(Duration)) * greatest(toUInt64OrZero(toString(SampleRate)), 1), toFloat64OrDefault(toString(Duration)) IS NOT NULL) / sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1), toFloat64OrDefault(toString(Duration)) IS NOT NULL) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite avg with where condition 1`] = `"SELECT sumIf(toFloat64OrDefault(toString(Duration)) * greatest(toUInt64OrZero(toString(SampleRate)), 1), ServiceName = 'api' AND toFloat64OrDefault(toString(Duration)) IS NOT NULL) / sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1), ServiceName = 'api' AND toFloat64OrDefault(toString(Duration)) IS NOT NULL) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) AND (ServiceName = 'api') SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite count() to sum(greatest(...)) 1`] = `"SELECT sum(greatest(toUInt64OrZero(toString(SampleRate)), 1)) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite countIf to sumIf(greatest(...), cond) 1`] = `"SELECT sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1), StatusCode = 'Error') FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) AND (StatusCode = 'Error') SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite quantile to quantileTDigestWeighted 1`] = `"SELECT quantileTDigestWeighted(0.99)(toFloat64OrDefault(toString(Duration)), toUInt32(greatest(toUInt64OrZero(toString(SampleRate)), 1))) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite quantile with where condition 1`] = `"SELECT quantileTDigestWeightedIf(0.95)(toFloat64OrDefault(toString(Duration)), toUInt32(greatest(toUInt64OrZero(toString(SampleRate)), 1)), ServiceName = 'api' AND toFloat64OrDefault(toString(Duration)) IS NOT NULL) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) AND (ServiceName = 'api') SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite sum to weighted sum 1`] = `"SELECT sum(toFloat64OrDefault(toString(Duration)) * greatest(toUInt64OrZero(toString(SampleRate)), 1)) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + +exports[`renderChartConfig sample-weighted aggregations should rewrite sum with where condition 1`] = `"SELECT sumIf(toFloat64OrDefault(toString(Duration)) * greatest(toUInt64OrZero(toString(SampleRate)), 1), ServiceName = 'api' AND toFloat64OrDefault(toString(Duration)) IS NOT NULL) FROM default.otel_traces WHERE (Timestamp >= fromUnixTimestamp64Milli(1739318400000) AND Timestamp <= fromUnixTimestamp64Milli(1739491200000)) AND (ServiceName = 'api') SETTINGS optimize_read_in_order = 0, cast_keep_nullable = 1, additional_result_filter = 'x != 2', count_distinct_implementation = 'uniqCombined64', async_insert_busy_timeout_min_ms = 20000"`; + exports[`renderChartConfig should generate sql for a single gauge metric 1`] = ` "WITH Source AS ( SELECT diff --git a/packages/common-utils/src/__tests__/renderChartConfig.test.ts b/packages/common-utils/src/__tests__/renderChartConfig.test.ts index 0d859fd701..9a57a313e6 100644 --- a/packages/common-utils/src/__tests__/renderChartConfig.test.ts +++ b/packages/common-utils/src/__tests__/renderChartConfig.test.ts @@ -1839,4 +1839,374 @@ describe('renderChartConfig', () => { ); }); }); + + describe('sample-weighted aggregations', () => { + const baseSampledConfig: ChartConfigWithOptDateRange = { + displayType: DisplayType.Table, + connection: 'test-connection', + from: { + databaseName: 'default', + tableName: 'otel_traces', + }, + select: [], + where: '', + whereLanguage: 'sql', + timestampValueExpression: 'Timestamp', + sampleWeightExpression: 'SampleRate', + dateRange: [new Date('2025-02-12'), new Date('2025-02-14')], + }; + + it('should rewrite count() to sum(greatest(...))', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'count', + valueExpression: '', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain( + 'greatest(toUInt64OrZero(toString(SampleRate)), 1)', + ); + expect(actual).toContain('sum('); + expect(actual).not.toContain('count()'); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite countIf to sumIf(greatest(...), cond)', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'count', + valueExpression: '', + aggCondition: "StatusCode = 'Error'", + aggConditionLanguage: 'sql', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain( + 'sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1)', + ); + expect(actual).not.toContain('countIf'); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite avg to weighted average', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'avg', + valueExpression: 'Duration', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain( + '* greatest(toUInt64OrZero(toString(SampleRate)), 1)', + ); + expect(actual).toContain( + '/ sumIf(greatest(toUInt64OrZero(toString(SampleRate)), 1)', + ); + expect(actual).not.toContain('avg('); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite sum to weighted sum', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'sum', + valueExpression: 'Duration', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain( + '* greatest(toUInt64OrZero(toString(SampleRate)), 1)', + ); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite quantile to quantileTDigestWeighted', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'quantile', + valueExpression: 'Duration', + aggCondition: '', + level: 0.99, + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('quantileTDigestWeighted(0.99)'); + expect(actual).toContain( + 'toUInt32(greatest(toUInt64OrZero(toString(SampleRate)), 1))', + ); + expect(actual).not.toContain('quantile(0.99)'); + expect(actual).toMatchSnapshot(); + }); + + it('should leave min/max unchanged with sampleWeightExpression', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'min', + valueExpression: 'Duration', + aggCondition: '', + }, + { + aggFn: 'max', + valueExpression: 'Duration', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('min('); + expect(actual).toContain('max('); + expect(actual).not.toContain('SampleRate'); + expect(actual).toMatchSnapshot(); + }); + + it('should leave count_distinct unchanged with sampleWeightExpression', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'count_distinct', + valueExpression: 'TraceId', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('count(DISTINCT'); + expect(actual).not.toContain('SampleRate'); + expect(actual).toMatchSnapshot(); + }); + + it('should handle complex sampleWeightExpression like SpanAttributes map access', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + sampleWeightExpression: "SpanAttributes['SampleRate']", + select: [ + { + aggFn: 'count', + valueExpression: '', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain( + "greatest(toUInt64OrZero(toString(SpanAttributes['SampleRate'])), 1)", + ); + expect(actual).toContain('sum('); + expect(actual).not.toContain('count()'); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite avg with where condition', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'avg', + valueExpression: 'Duration', + aggCondition: "ServiceName = 'api'", + aggConditionLanguage: 'sql', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('sumIf('); + expect(actual).toContain("ServiceName = 'api'"); + expect(actual).not.toContain('avg('); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite sum with where condition', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'sum', + valueExpression: 'Duration', + aggCondition: "ServiceName = 'api'", + aggConditionLanguage: 'sql', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('sumIf('); + expect(actual).toContain("ServiceName = 'api'"); + expect(actual).toMatchSnapshot(); + }); + + it('should rewrite quantile with where condition', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'quantile', + valueExpression: 'Duration', + aggCondition: "ServiceName = 'api'", + aggConditionLanguage: 'sql', + level: 0.95, + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('quantileTDigestWeightedIf(0.95)'); + expect(actual).toContain("ServiceName = 'api'"); + expect(actual).not.toContain('quantile(0.95)'); + expect(actual).toMatchSnapshot(); + }); + + it('should handle mixed weighted and passthrough aggregations', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + select: [ + { + aggFn: 'count', + valueExpression: '', + aggCondition: '', + alias: 'weighted_count', + }, + { + aggFn: 'avg', + valueExpression: 'Duration', + aggCondition: '', + alias: 'weighted_avg', + }, + { + aggFn: 'min', + valueExpression: 'Duration', + aggCondition: '', + alias: 'min_duration', + }, + { + aggFn: 'count_distinct', + valueExpression: 'TraceId', + aggCondition: '', + alias: 'unique_traces', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('sum('); + expect(actual).toContain('min('); + expect(actual).toContain('count(DISTINCT'); + expect(actual).not.toContain('count()'); + expect(actual).not.toContain('avg('); + expect(actual).toMatchSnapshot(); + }); + + it('should not rewrite aggregations without sampleWeightExpression', async () => { + const config: ChartConfigWithOptDateRange = { + ...baseSampledConfig, + sampleWeightExpression: undefined, + select: [ + { + aggFn: 'count', + valueExpression: '', + aggCondition: '', + }, + ], + }; + + const generatedSql = await renderChartConfig( + config, + mockMetadata, + querySettings, + ); + const actual = parameterizedQueryToSql(generatedSql); + expect(actual).toContain('count()'); + expect(actual).not.toContain('SampleRate'); + }); + }); }); diff --git a/packages/common-utils/src/core/renderChartConfig.ts b/packages/common-utils/src/core/renderChartConfig.ts index 0707646700..3149d3e925 100644 --- a/packages/common-utils/src/core/renderChartConfig.ts +++ b/packages/common-utils/src/core/renderChartConfig.ts @@ -319,11 +319,13 @@ const aggFnExpr = ({ expr, level, where, + sampleWeightExpression, }: { fn: AggregateFunction | AggregateFunctionWithCombinators; expr?: string; level?: number; where?: string; + sampleWeightExpression?: string; }) => { const isAny = fn === 'any'; const isNone = fn === 'none'; @@ -365,6 +367,79 @@ const aggFnExpr = ({ })`; } + // Sample-weighted aggregations: when sampleWeightExpression is set, + // each row carries a weight (defaults to 1 for unsampled spans). + // Corrected formulas account for upstream sampling (1-in-N). + // The greatest(..., 1) ensures unsampled rows (missing/empty/zero) + // are counted at weight 1 rather than dropped. + if ( + sampleWeightExpression && + !fn.endsWith('Merge') && + !fn.endsWith('State') + ) { + const sampleWeightExpr = `greatest(toUInt64OrZero(toString(${sampleWeightExpression})), 1)`; + const w = { UNSAFE_RAW_SQL: sampleWeightExpr }; + + if (fn === 'count') { + return isWhereUsed + ? chSql`sumIf(${w}, ${{ UNSAFE_RAW_SQL: where }})` + : chSql`sum(${w})`; + } + + if (fn === 'none') { + return chSql`${{ UNSAFE_RAW_SQL: expr ?? '' }}`; + } + + if (expr != null) { + if (fn === 'count_distinct' || fn === 'min' || fn === 'max') { + // These cannot be corrected for sampling; pass through unchanged + if (fn === 'count_distinct') { + return chSql`count${isWhereUsed ? 'If' : ''}(DISTINCT ${{ + UNSAFE_RAW_SQL: expr, + }}${isWhereUsed ? chSql`, ${{ UNSAFE_RAW_SQL: where }}` : ''})`; + } + return chSql`${{ UNSAFE_RAW_SQL: fn }}${isWhereUsed ? 'If' : ''}( + ${unsafeExpr}${isWhereUsed ? chSql`, ${{ UNSAFE_RAW_SQL: whereWithExtraNullCheck }}` : ''} + )`; + } + + if (fn === 'avg') { + const weightedVal = { + UNSAFE_RAW_SQL: `${unsafeExpr.UNSAFE_RAW_SQL} * ${sampleWeightExpr}`, + }; + const nullCheck = `${unsafeExpr.UNSAFE_RAW_SQL} IS NOT NULL`; + if (isWhereUsed) { + const cond = { UNSAFE_RAW_SQL: `${where} AND ${nullCheck}` }; + return chSql`sumIf(${weightedVal}, ${cond}) / sumIf(${w}, ${cond})`; + } + return chSql`sumIf(${weightedVal}, ${{ UNSAFE_RAW_SQL: nullCheck }}) / sumIf(${w}, ${{ UNSAFE_RAW_SQL: nullCheck }})`; + } + + if (fn === 'sum') { + const weightedVal = { + UNSAFE_RAW_SQL: `${unsafeExpr.UNSAFE_RAW_SQL} * ${sampleWeightExpr}`, + }; + if (isWhereUsed) { + return chSql`sumIf(${weightedVal}, ${{ UNSAFE_RAW_SQL: whereWithExtraNullCheck }})`; + } + return chSql`sum(${weightedVal})`; + } + + if (level != null && fn.startsWith('quantile')) { + const levelStr = Number.isFinite(level) ? `${level}` : '0'; + const weightArg = { + UNSAFE_RAW_SQL: `toUInt32(${sampleWeightExpr})`, + }; + if (isWhereUsed) { + return chSql`quantileTDigestWeightedIf(${{ UNSAFE_RAW_SQL: levelStr }})(${unsafeExpr}, ${weightArg}, ${{ UNSAFE_RAW_SQL: whereWithExtraNullCheck }})`; + } + return chSql`quantileTDigestWeighted(${{ UNSAFE_RAW_SQL: levelStr }})(${unsafeExpr}, ${weightArg})`; + } + + // For any other fn (last_value, any, etc.), fall through to default + } + } + if (fn === 'count') { if (isWhereUsed) { return chSql`${fn}If(${{ UNSAFE_RAW_SQL: where }})`; @@ -477,12 +552,14 @@ async function renderSelectList( // @ts-expect-error (TS doesn't know that we've already checked for quantile) level: select.level, where: whereClause.sql, + sampleWeightExpression: chartConfig.sampleWeightExpression, }); } else { expr = aggFnExpr({ fn: select.aggFn, expr: select.valueExpression, where: whereClause.sql, + sampleWeightExpression: chartConfig.sampleWeightExpression, }); } diff --git a/packages/common-utils/src/types.ts b/packages/common-utils/src/types.ts index 42d568d3c7..61659fcaf9 100644 --- a/packages/common-utils/src/types.ts +++ b/packages/common-utils/src/types.ts @@ -520,6 +520,7 @@ const SharedChartDisplaySettingsSchema = z.object({ export const _ChartConfigSchema = SharedChartDisplaySettingsSchema.extend({ timestampValueExpression: z.string(), implicitColumnExpression: z.string().optional(), + sampleWeightExpression: z.string().optional(), markdown: z.string().optional(), filtersLogicalOperator: z.enum(['AND', 'OR']).optional(), filters: z.array(FilterSchema).optional(), @@ -927,6 +928,7 @@ export const TraceSourceSchema = BaseSourceSchema.extend({ spanKindExpression: z.string().min(1, 'Span Kind Expression is required'), // Optional fields for traces + sampleRateExpression: z.string().optional(), logSourceId: z.string().optional().nullable(), sessionSourceId: z.string().optional(), metricSourceId: z.string().optional(),