Aggregation improvements for the legend

Author: matt-aitken

apps/webapp/app/components/code/QueryResultsChart.tsx

@@ -3,7 +3,8 @@ import { memo, useMemo } from "react";
 import type { ChartConfig } from "~/components/primitives/charts/Chart";
 import { Chart } from "~/components/primitives/charts/ChartCompound";
 import { Paragraph } from "../primitives/Paragraph";
-import { AggregationType, ChartConfiguration } from "../metrics/QueryWidget";
+import type { AggregationType, ChartConfiguration } from "../metrics/QueryWidget";
+import { aggregateValues } from "../primitives/charts/aggregation";
 import { getRunStatusHexColor } from "~/components/runs/v3/TaskRunStatus";
 import { getSeriesColor } from "./chartColors";
 
@@ -671,25 +672,6 @@ function toNumber(value: unknown): number {
   return 0;
 }
 
-/**
- * Aggregate an array of numbers using the specified aggregation function
- */
-function aggregateValues(values: number[], aggregation: AggregationType): number {
-  if (values.length === 0) return 0;
-  switch (aggregation) {
-    case "sum":
-      return values.reduce((a, b) => a + b, 0);
-    case "avg":
-      return values.reduce((a, b) => a + b, 0) / values.length;
-    case "count":
-      return values.length;
-    case "min":
-      return Math.min(...values);
-    case "max":
-      return Math.max(...values);
-  }
-}
-
 /**
  * Sort data array by a specified column
  */
@@ -1032,6 +1014,7 @@ export const QueryResultsChart = memo(function QueryResultsChart({
         labelFormatter={legendLabelFormatter}
         showLegend={showLegend}
         maxLegendItems={fullLegend ? Infinity : 5}
+        legendAggregation={config.aggregation}
         minHeight="300px"
         fillContainer
         onViewAllLegendItems={onViewAllLegendItems}
@@ -1058,6 +1041,7 @@ export const QueryResultsChart = memo(function QueryResultsChart({
       labelFormatter={legendLabelFormatter}
       showLegend={showLegend}
       maxLegendItems={fullLegend ? Infinity : 5}
+      legendAggregation={config.aggregation}
       minHeight="300px"
       fillContainer
       onViewAllLegendItems={onViewAllLegendItems}

apps/webapp/app/components/primitives/charts/ChartLegendCompound.tsx

@@ -1,18 +1,30 @@
 import React, { useMemo } from "react";
+import type { AggregationType } from "~/components/metrics/QueryWidget";
 import { useChartContext } from "./ChartContext";
 import { useSeriesTotal } from "./ChartRoot";
+import { aggregateValues } from "./aggregation";
 import { cn } from "~/utils/cn";
 import { AnimatedNumber } from "../AnimatedNumber";
 
+const aggregationLabels: Record<AggregationType, string> = {
+  sum: "Sum",
+  avg: "Average",
+  count: "Count",
+  min: "Min",
+  max: "Max",
+};
+
 export type ChartLegendCompoundProps = {
   /** Maximum number of legend items to show before collapsing */
   maxItems?: number;
   /** Hide the legend entirely (useful for conditional rendering) */
   hidden?: boolean;
   /** Additional className */
   className?: string;
-  /** Label for the total row */
+  /** Label for the total row (derived from aggregation when not provided) */
   totalLabel?: string;
+  /** Aggregation method – controls the header label and how totals are computed */
+  aggregation?: AggregationType;
   /** Callback when "View all" button is clicked */
   onViewAllLegendItems?: () => void;
   /** When true, constrains legend to max 50% height with scrolling */
@@ -35,45 +47,59 @@ export function ChartLegendCompound({
   maxItems = Infinity,
   hidden = false,
   className,
-  totalLabel = "Total",
+  totalLabel,
+  aggregation,
   onViewAllLegendItems,
   scrollable = false,
 }: ChartLegendCompoundProps) {
   const { config, dataKey, dataKeys, highlight, labelFormatter } = useChartContext();
-  const totals = useSeriesTotal();
+  const totals = useSeriesTotal(aggregation);
 
-  // Calculate grand total (sum of all series totals)
+  // Derive the effective label from the aggregation type when no explicit label is provided
+  const effectiveTotalLabel = totalLabel ?? (aggregation ? aggregationLabels[aggregation] : "Total");
+
+  // Calculate grand total by aggregating across all per-series values
   const grandTotal = useMemo(() => {
-    return dataKeys.reduce((sum, key) => sum + (totals[key] || 0), 0);
-  }, [totals, dataKeys]);
+    const values = dataKeys.map((key) => totals[key] || 0);
+    if (!aggregation) {
+      // Default: sum
+      return values.reduce((a, b) => a + b, 0);
+    }
+    return aggregateValues(values, aggregation);
+  }, [totals, dataKeys, aggregation]);
 
   // Calculate current total based on hover state
   const currentTotal = useMemo(() => {
     if (!highlight.activePayload?.length) return grandTotal;
 
-    // Sum all values from the hovered data point
-    return highlight.activePayload.reduce((sum, item) => {
-      if (item.value !== undefined && dataKeys.includes(item.dataKey as string)) {
-        return sum + (Number(item.value) || 0);
-      }
-      return sum;
-    }, 0);
-  }, [highlight.activePayload, grandTotal, dataKeys]);
+    // Collect all series values from the hovered data point
+    const values = highlight.activePayload
+      .filter((item) => item.value !== undefined && dataKeys.includes(item.dataKey as string))
+      .map((item) => Number(item.value) || 0);
+
+    if (values.length === 0) return 0;
+
+    if (!aggregation) {
+      // Default: sum
+      return values.reduce((a, b) => a + b, 0);
+    }
+    return aggregateValues(values, aggregation);
+  }, [highlight.activePayload, grandTotal, dataKeys, aggregation]);
 
-  // Get the label for the total row - x-axis value when hovering, totalLabel otherwise
+  // Get the label for the total row - x-axis value when hovering, effectiveTotalLabel otherwise
   const currentTotalLabel = useMemo(() => {
-    if (!highlight.activePayload?.length) return totalLabel;
+    if (!highlight.activePayload?.length) return effectiveTotalLabel;
 
     // Get the x-axis label from the payload's original data
     const firstPayloadItem = highlight.activePayload[0];
     const xAxisValue = firstPayloadItem?.payload?.[dataKey];
 
-    if (xAxisValue === undefined) return totalLabel;
+    if (xAxisValue === undefined) return effectiveTotalLabel;
 
     // Apply the formatter if provided, otherwise just stringify the value
     const stringValue = String(xAxisValue);
     return labelFormatter ? labelFormatter(stringValue) : stringValue;
-  }, [highlight.activePayload, dataKey, totalLabel, labelFormatter]);
+  }, [highlight.activePayload, dataKey, effectiveTotalLabel, labelFormatter]);
 
   // Get current data for the legend based on hover state
   const currentData = useMemo(() => {

apps/webapp/app/components/primitives/charts/ChartRoot.tsx

@@ -1,5 +1,6 @@
 import React, { useMemo } from "react";
 import type * as RechartsPrimitive from "recharts";
+import type { AggregationType } from "~/components/metrics/QueryWidget";
 import { ChartContainer, type ChartConfig, type ChartState } from "./Chart";
 import { ChartProvider, useChartContext, type LabelFormatter } from "./ChartContext";
 import { ChartLegendCompound } from "./ChartLegendCompound";
@@ -29,6 +30,8 @@ export type ChartRootProps = {
   maxLegendItems?: number;
   /** Label for the total row in the legend */
   legendTotalLabel?: string;
+  /** Aggregation method used by the legend to compute totals (defaults to sum behavior) */
+  legendAggregation?: AggregationType;
   /** Callback when "View all" legend button is clicked */
   onViewAllLegendItems?: () => void;
   /** When true, constrains legend to max 50% height with scrolling */
@@ -73,6 +76,7 @@ export function ChartRoot({
   showLegend = false,
   maxLegendItems = 5,
   legendTotalLabel,
+  legendAggregation,
   onViewAllLegendItems,
   legendScrollable = false,
   fillContainer = false,
@@ -96,6 +100,7 @@ export function ChartRoot({
         showLegend={showLegend}
         maxLegendItems={maxLegendItems}
         legendTotalLabel={legendTotalLabel}
+        legendAggregation={legendAggregation}
         onViewAllLegendItems={onViewAllLegendItems}
         legendScrollable={legendScrollable}
         fillContainer={fillContainer}
@@ -112,6 +117,7 @@ type ChartRootInnerProps = {
   showLegend?: boolean;
   maxLegendItems?: number;
   legendTotalLabel?: string;
+  legendAggregation?: AggregationType;
   onViewAllLegendItems?: () => void;
   legendScrollable?: boolean;
   fillContainer?: boolean;
@@ -124,6 +130,7 @@ function ChartRootInner({
   showLegend = false,
   maxLegendItems = 5,
   legendTotalLabel,
+  legendAggregation,
   onViewAllLegendItems,
   legendScrollable = false,
   fillContainer = false,
@@ -165,6 +172,7 @@ function ChartRootInner({
         <ChartLegendCompound
           maxItems={maxLegendItems}
           totalLabel={legendTotalLabel}
+          aggregation={legendAggregation}
           onViewAllLegendItems={onViewAllLegendItems}
           scrollable={legendScrollable}
         />
@@ -194,18 +202,75 @@ export function useHasNoData(): boolean {
 }
 
 /**
- * Hook to calculate totals for each series across all data points.
+ * Hook to calculate aggregated values for each series across all data points.
+ * When no aggregation is provided, defaults to sum (original behavior).
  * Useful for legend displays.
  */
-export function useSeriesTotal(): Record<string, number> {
+export function useSeriesTotal(aggregation?: AggregationType): Record<string, number> {
   const { data, dataKeys } = useChartContext();
 
   return useMemo(() => {
-    return data.reduce((acc, item) => {
+    // Sum (default) and count use additive accumulation
+    if (!aggregation || aggregation === "sum" || aggregation === "count") {
+      return data.reduce(
+        (acc, item) => {
+          for (const seriesKey of dataKeys) {
+            acc[seriesKey] = (acc[seriesKey] || 0) + Number(item[seriesKey] || 0);
+          }
+          return acc;
+        },
+        {} as Record<string, number>
+      );
+    }
+
+    if (aggregation === "avg") {
+      const sums: Record<string, number> = {};
+      const counts: Record<string, number> = {};
+      for (const item of data) {
+        for (const seriesKey of dataKeys) {
+          const val = Number(item[seriesKey] || 0);
+          sums[seriesKey] = (sums[seriesKey] || 0) + val;
+          counts[seriesKey] = (counts[seriesKey] || 0) + 1;
+        }
+      }
+      const result: Record<string, number> = {};
+      for (const key of dataKeys) {
+        result[key] = counts[key] ? sums[key]! / counts[key]! : 0;
+      }
+      return result;
+    }
+
+    if (aggregation === "min") {
+      const result: Record<string, number> = {};
+      for (const item of data) {
+        for (const seriesKey of dataKeys) {
+          const val = Number(item[seriesKey] || 0);
+          if (result[seriesKey] === undefined || val < result[seriesKey]) {
+            result[seriesKey] = val;
+          }
+        }
+      }
+      // Default to 0 for series with no data
+      for (const key of dataKeys) {
+        if (result[key] === undefined) result[key] = 0;
+      }
+      return result;
+    }
+
+    // aggregation === "max"
+    const result: Record<string, number> = {};
+    for (const item of data) {
       for (const seriesKey of dataKeys) {
-        acc[seriesKey] = (acc[seriesKey] || 0) + Number(item[seriesKey] || 0);
+        const val = Number(item[seriesKey] || 0);
+        if (result[seriesKey] === undefined || val > result[seriesKey]) {
+          result[seriesKey] = val;
+        }
       }
-      return acc;
-    }, {} as Record<string, number>);
-  }, [data, dataKeys]);
+    }
+    // Default to 0 for series with no data
+    for (const key of dataKeys) {
+      if (result[key] === undefined) result[key] = 0;
+    }
+    return result;
+  }, [data, dataKeys, aggregation]);
 }

apps/webapp/app/components/primitives/charts/aggregation.ts

@@ -0,0 +1,23 @@
+import type { AggregationType } from "~/components/metrics/QueryWidget";
+
+/**
+ * Aggregate an array of numbers using the specified aggregation function.
+ *
+ * Shared utility so both QueryResultsChart (data transformation) and chart
+ * legend components can reuse the same logic without circular imports.
+ */
+export function aggregateValues(values: number[], aggregation: AggregationType): number {
+  if (values.length === 0) return 0;
+  switch (aggregation) {
+    case "sum":
+      return values.reduce((a, b) => a + b, 0);
+    case "avg":
+      return values.reduce((a, b) => a + b, 0) / values.length;
+    case "count":
+      return values.length;
+    case "min":
+      return Math.min(...values);
+    case "max":
+      return Math.max(...values);
+  }
+}